fix: mapreduce with unitrange dims (#1572)

Author: avik-pal

Project.toml

@@ -1,7 +1,7 @@
 name = "Reactant"
 uuid = "3c362404-f566-11ee-1572-e11a4b42c853"
 authors = ["William Moses <[email protected]>", "Valentin Churavy <[email protected]>", "Sergio Sánchez Ramírez <[email protected]>", "Paul Berg <[email protected]>", "Avik Pal <[email protected]>", "Mosè Giordano <[email protected]>"]
-version = "0.2.154"
+version = "0.2.155"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/TracedRArray.jl

@@ -568,7 +568,7 @@ function overloaded_mapreduce(
     original_dims = dims
     dims isa Int && (dims = Int64[dims])
     dims isa Colon && (dims = collect(Int64, 1:N))
-    dims isa AbstractVector{<:Integer} || (dims = collect(Int64, dims))
+    dims isa Vector{Int64} || (dims = collect(Int64, dims))
 
     op_in_T = unwrapped_eltype(Core.Compiler.return_type(f, Tuple{T}))
     reduce_init = __default_init(op_in_T, op)

test/basic.jl

@@ -1524,3 +1524,10 @@ end
         )
     end
 end
+
+@testset "mapreduce with unitrange dims" begin
+    x = reshape(collect(Float32, 1:64), 2, 4, 8)
+    x_ra = Reactant.to_rarray(x)
+
+    @test @jit(sum(x_ra; dims=1:2)) ≈ sum(x; dims=1:2)
+end