feat: split out non-generator changes from #1642

Author: avik-pal

src/Reactant.jl

@@ -187,6 +187,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

@@ -1500,29 +1500,53 @@ 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...) = Reactant.call_with_reactant(fn.f, (args...,))
+(fn::BroadcastIterator)(args...) = fn.f((args...,))
 
 function unwrapped_broadcast(f::F, x::Base.Iterators.Zip) where {F}
     min_length = Base.inferencebarrier(minimum)(length, x.is)
     itrs = [length(itr) > min_length ? itr[1:min_length] : itr for itr in x.is]
     if any(Base.Fix2(isa, AnyTracedRArray), itrs)
-        return (BroadcastIterator(f)).(itrs...)
+        return broadcast(BroadcastIterator(f), itrs...)
     else
-        fn = BroadcastIterator(f)
-        return [fn(Base.Fix2(getindex, i).(itrs)...) for i in 1:min_length]
+        return unwrapped_broadcast_with_iterate(f, x)
     end
 end
 
 function unwrapped_broadcast(f::F, x::Base.Iterators.Enumerate) where {F}
     if x.itr isa AnyTracedRArray
-        return (BroadcastIterator(f)).(1:length(x.itr), x.itr)
+        return broadcast(
+            BroadcastIterator(f), Reactant.promote_to(TracedRArray, 1:length(x.itr)), x.itr
+        )
     else
-        return [f((i, x.itr[i])) for i in 1:length(x.itr)]
+        return unwrapped_broadcast_with_iterate(f, x)
     end
 end
 
-unwrapped_broadcast(f::F, xs::Vector) where {F} = [f(x) for x in xs]
+unwrapped_broadcast(f::F, xs) where {F} = unwrapped_broadcast_with_iterate(f, xs)
+
+function unwrapped_broadcast_with_iterate(f::F, itr) where {F}
+    y = Reactant.call_with_reactant(iterate, itr)
+    y === nothing && return []
+
+    first, state = y
+    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 = @opcall call(f, val)
+        push!(result, res)
+    end
+
+    return result
+end
 
 end

src/utils.jl

@@ -1,4 +1,4 @@
-struct CallWithReactant{F}
+struct CallWithReactant{F} <: Function
     f::F
 end
 

test/autodiff.jl

@@ -132,7 +132,7 @@ end
 
 @testset "Forward Gradient" begin
     x = Reactant.to_rarray(3.1 * ones(2, 2))
-    res = @test_warn r"`Adapt.parent_type` is not implemented for" @jit gw(x)
+    res = @jit gw(x)
     # TODO we should probably override https://github.com/EnzymeAD/Enzyme.jl/blob/5e6a82dd08e74666822b9d7b2b46c36b075668ca/src/Enzyme.jl#L2132
     # to make sure this gets merged as a tracedrarray
     @test res isa Tuple{<:Enzyme.TupleArray{<:ConcreteRNumber{Float64},(2, 2),4,2}}

test/basic.jl

@@ -926,7 +926,7 @@ end
 
     ra = Reactant.to_rarray(x)
     @jit dip!(ra)
-    ra[:a] ≈ (2.7 * 2) * ones(4)
+    @test ra[:a] ≈ (2.7 * 3.1) * ones(4)
 end
 
 @testset "@code_xla" begin
@@ -1429,7 +1429,10 @@ end
 end
 
 zip_iterator(a, b) = mapreduce(splat(*), +, zip(a, b))
+zip_iterator2(a, b) = mapreduce(splat(.-), +, zip(a, b))
 enumerate_iterator(a) = mapreduce(splat(*), +, enumerate(a))
+enumerate_iterator2(a) = mapreduce(splat(.-), +, enumerate(a))
+mapreduce_vector(a) = mapreduce(-, +, a)
 
 function nested_mapreduce_zip(x, y)
     return mapreduce(+, zip(eachcol(x), eachcol(y)); init=0.0f0) do (x, y)
@@ -1448,18 +1451,30 @@ end
 @testset "Base.Iterators" begin
     @testset "zip" begin
         N = 10
-        a = range(1.0, 5.0; length=N)
-        x = range(10.0, 15.0; length=N + 2)
+        a = collect(range(1.0, 5.0; length=N))
+        x = collect(range(10.0, 15.0; length=N + 2))
         x_ra = Reactant.to_rarray(x)
 
         @test @jit(zip_iterator(a, x_ra)) ≈ zip_iterator(a, x)
+
+        a = [rand(Float32, 2, 3) for _ in 1:10]
+        x = [rand(Float32, 2, 3) for _ in 1:10]
+        a_ra = Reactant.to_rarray(a)
+        x_ra = Reactant.to_rarray(x)
+
+        @test @jit(zip_iterator2(a_ra, x_ra)) ≈ zip_iterator2(a, x)
     end
 
     @testset "enumerate" begin
-        x = range(1.0, 5.0; length=10)
+        x = collect(range(1.0, 5.0; length=10))
         x_ra = Reactant.to_rarray(x)
 
         @test @jit(enumerate_iterator(x_ra)) ≈ enumerate_iterator(x)
+
+        x = [rand(Float32, 2, 3) for _ in 1:10]
+        x_ra = Reactant.to_rarray(x)
+
+        @test @jit(enumerate_iterator2(x_ra)) ≈ enumerate_iterator2(x)
     end
 
     @testset "nested mapreduce" begin
@@ -1481,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 "compilation cache" begin