fix: correct dims handling in mapreducedim! (#728)

* feat: add sign dispatches

* fix: correct dims handling in mapreducedim!

* Update test/basic.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Author: avik-pal

src/TracedRArray.jl

@@ -529,10 +529,14 @@ function Base.mapreducedim!(
     @nospecialize(R::TracedRArray),
     A::Base.AbstractArrayOrBroadcasted,
 )
-    tmp = TracedUtils.broadcast_to_size(
-        Base.mapreduce(f, op, A; dims=1), (1, size(R)[2:end]...)
-    )
-    R.mlir_data = broadcast(op, R, tmp).mlir_data
+    @assert length(size(R)) == length(size(A))
+    dims = map(enumerate(zip(size(R), size(A)))) do (i, (sR, sA))
+        sR == sA && return nothing
+        @assert sR == 1
+        return i
+    end
+    tmp = mapreduce(f, op, A; dims=filter(!isnothing, dims))
+    set_mlir_data!(R, get_mlir_data(tmp))
     return R
 end
 

test/basic.jl

@@ -850,3 +850,30 @@ end
 
     @test @jit(fn(x)) ≈ fn(Array(x))
 end
+
+function fntest1(x)
+    y = similar(x, 1, 1, 8)
+    sum!(y, x)
+    return y
+end
+
+function fntest2(x)
+    y = similar(x, 2, 1, 8)
+    sum!(y, x)
+    return y
+end
+
+function fntest3(x)
+    y = similar(x, 2, 1, 1)
+    sum!(abs2, y, x)
+    return y
+end
+
+@testset "mapreducedim!" begin
+    x = reshape(collect(Float32, 1:64), 2, 4, 8) ./ 64
+    x_ra = Reactant.to_rarray(x)
+
+    @test Array(@jit(fntest1(x_ra))) ≈ fntest1(x)
+    @test Array(@jit(fntest2(x_ra))) ≈ fntest2(x)
+    @test Array(@jit(fntest3(x_ra))) ≈ fntest3(x)
+end