fix: drop=false

Author: avik-pal

src/TracedRArray.jl

@@ -1402,26 +1402,27 @@ function unwrapped_broadcast(f::F, x::Slices, original_dims) where {F}
         updated_dims = ()
         if original_dims isa Colon
             updated_dims = mapslices_dims
+            re = x -> dropdims(x; dims=mapslices_dims)
         else
             for d in original_dims
                 idx = findfirst(isequal(d), x.slicemap)
                 @assert idx !== nothing "Expected dimension $d in $(x.slicemap)"
                 updated_dims = (updated_dims..., idx)
             end
+            re = x -> eachslice(x; dims=mapslices_dims, drop=true)
         end
 
-        return (
-            mapslices(f, px; dims=mapslices_dims),
-            updated_dims,
-            x -> eachslice(x; dims=mapslices_dims, drop=true),
-        )
+        return mapslices(f, px; dims=mapslices_dims), updated_dims, re
     else
         mapslices_dims = Tuple(filter(i -> !(x.slicemap[i] isa Colon), 1:ndims(px)))
-        return (
-            mapslices(f, px; dims=mapslices_dims),
-            original_dims,
-            x -> eachslice(x; dims=mapslices_dims, drop=false),
-        )
+        if original_dims isa Colon
+            updated_dims = mapslices_dims
+            re = x -> dropdims(x; dims=mapslices_dims)
+        else
+            updated_dims = Tuple(d for d in original_dims if d in mapslices_dims)
+            re = x -> eachslice(x; dims=mapslices_dims, drop=false)
+        end
+        return mapslices(f, px; dims=mapslices_dims), updated_dims, re
     end
 end
 

test/basic.jl

@@ -1742,3 +1742,43 @@ end
     fn = @compile sum(x_ra1)
     @test_throws Reactant.Compiler.MisMatchedThunkTypeError fn(x_ra2)
 end
+
+@testset "Slices" begin
+    @testset "drop=true" begin
+        x = eachslice(
+            Reactant.TestUtils.construct_test_array(Float32, 2, 3, 4, 5); dims=(3, 1)
+        )
+        x_ra = Reactant.to_rarray(x)
+
+        @test @jit(sum(x_ra)) ≈ sum(x)
+
+        @testset for dims in (1, 2, (1, 2), (2, 1))
+            res_ra = @jit sum(x_ra; dims)
+            res = sum(x; dims)
+            @test size(res_ra) == size(res)
+            for (gt, comp) in zip(res_ra, res)
+                @test gt ≈ comp
+            end
+        end
+    end
+
+    @testset "drop=false" begin
+        x = eachslice(
+            Reactant.TestUtils.construct_test_array(Float32, 2, 3, 4, 5);
+            dims=(3, 1),
+            drop=false,
+        )
+        x_ra = Reactant.to_rarray(x)
+
+        @test @jit(sum(x_ra)) ≈ sum(x)
+
+        @testset for dims in (1, 2, 3, 4, (1, 2), (1, 2, 4), (3, 4, 1), (2, 1))
+            res_ra = @jit sum(x_ra; dims)
+            res = sum(x; dims)
+            @test size(res_ra) == size(res)
+            for (gt, comp) in zip(res_ra, res)
+                @test gt ≈ comp
+            end
+        end
+    end
+end