fix: correct semantics for Colon mapreduce

Author: avik-pal

src/Compiler.jl

@@ -575,7 +575,7 @@ end
 function compile_xla(f, args; client=nothing)
     # register MLIR dialects
     ctx = MLIR.IR.Context()
-    Base.append!(Reactant.registry[]; context=ctx)
+    append!(Reactant.registry[]; context=ctx)
     @ccall MLIR.API.mlir_c.RegisterDialects(ctx::MLIR.API.MlirContext)::Cvoid
 
     return MLIR.IR.context!(ctx) do

src/TracedRArray.jl

@@ -431,7 +431,11 @@ function Base.mapreduce(
         )
         red = TracedRArray{T,length(toonedims)}((), red, (toonedims...,))
     else
-        red = TracedRArray{T,length(outdims)}((), red, (outdims...,))
+        if length(outdims) == 0
+            red = TracedRNumber{T}((), red)
+        else
+            red = TracedRArray{T,length(outdims)}((), red, (outdims...,))
+        end
     end
     return red
 end

test/basic.jl

@@ -65,6 +65,8 @@ end
 
 sumexp(x) = sum(exp, x)
 
+sum_compare(x) = sum(x) > 0
+
 @testset "Basic mapreduce" begin
     x = rand(Float32, 10)
     a = Reactant.ConcreteRArray(x)
@@ -74,6 +76,12 @@ sumexp(x) = sum(exp, x)
     f_res = f(a)
 
     @test f_res ≈ r_res
+
+    # Ensure we are tracing as scalars. Else this will fail due to > not being defined on
+    # arrays
+    f = @compile sum_compare(a)
+    # We need to use [] to unwrap the scalar. We will fix this in the future.
+    @test f(a)[] == sum_compare(x)
 end
 
 function mysoftmax!(x)