feat: split out non-generator changes from #1642

Author: avik-pal

src/Reactant.jl

@@ -193,6 +193,7 @@ use_overlayed_version(::TracedRArray) = true
 use_overlayed_version(::TracedRNumber) = true
 use_overlayed_version(::Number) = false
 use_overlayed_version(::MissingTracedValue) = true
+use_overlayed_version(::Vector{<:AnyTracedRArray}) = true
 use_overlayed_version(::AbstractArray{<:TracedRNumber}) = true
 use_overlayed_version(rng::ReactantRNG) = use_overlayed_version(rng.seed)
 function use_overlayed_version(x::AbstractArray)

src/TracedRArray.jl

@@ -1509,6 +1509,9 @@ end
 
 struct BroadcastIterator{F}
     f::F
+
+    BroadcastIterator{F}(f::F) where {F} = new{F}(f)
+    BroadcastIterator(f::F) where {F} = new{F}(f)
 end
 
 (fn::BroadcastIterator)(args...) = fn.f((args...,))
@@ -1536,15 +1539,15 @@ function unrolled_map(f::F, itr) where {F}
     y === nothing && return []
 
     first, state = y
-    res_first = Reactant.call_with_reactant(f, first)
+    res_first = @opcall call(f, first)
     result = [res_first]
 
     while true
         y = Reactant.call_with_reactant(iterate, itr, state)
         y === nothing && break
 
         val, state = y
-        res = Reactant.call_with_reactant(f, val)
+        res = @opcall call(f, val)
         push!(result, res)
     end
 

test/basic.jl

@@ -1496,6 +1496,15 @@ end
 
         @test @jit(nested_mapreduce_hcat(x_ra, y_ra)) ≈ nested_mapreduce_hcat(x, y)
     end
+
+    @testset "mapreduce vector" begin
+        x = [rand(Float32, 2, 3) for _ in 1:10]
+        x_ra = Reactant.to_rarray(x)
+
+        @test @jit(mapreduce_vector(x_ra)) ≈ mapreduce_vector(x)
+        hlo = repr(@code_hlo optimize = false mapreduce_vector(x_ra))
+        @test contains(hlo, "call")
+    end
 end
 
 @testset "Base.Generator" begin

test/integration/fillarrays.jl

@@ -25,5 +25,5 @@ end
     x = OneElement(3.4f0, (3, 4), (32, 32))
     rx = Reactant.to_rarray(x)
 
-    @test @jit(fn(rx, rx)) ≈ fn(x, x)
+    @test @jit(fn(rx, rx)) ≈ fn(x, x) atol = 1e-3 rtol = 1e-3
 end