fix: setindex with traced indices (#1197)

Author: avik-pal

src/TracedRArray.jl

@@ -377,7 +377,9 @@ function Base.setindex!(a::TracedRArray{T,N}, v, indices::Vararg{Any,N}) where {
         v = TracedUtils.promote_to(TracedRArray{T,N}, v)
     else
         v = TracedUtils.promote_to(TracedRArray{T,ndims(v)}, v)
-        non_integer_indices = [!(idx isa Integer) for idx in indices]
+        non_integer_indices = [
+            !(idx isa Union{Integer,TracedRNumber{<:Integer}}) for idx in indices
+        ]
         broadcast_dims = findall(non_integer_indices)
         if length(broadcast_dims) == N
             v = TracedUtils.broadcast_to_size(v, length.(indices))

test/control_flow.jl

@@ -836,3 +836,19 @@ end
     recursivenottraced.x[3].x = recursivenottraced
     @test !Reactant.ReactantCore.is_traced(recursivenottraced)
 end
+
+function loop_batched(x)
+    y = similar(x)
+    @trace for i in 1:size(x, 1)
+        y[i, :] = x[i, :] .+ 1
+        y[i, :] = y[i, :] .^ 2
+    end
+    return y
+end
+
+@testset "setindex: batched" begin
+    x = rand(1024, 128)
+    x_ra = Reactant.to_rarray(x)
+
+    @test @jit(loop_batched(x_ra)) ≈ loop_batched(x)
+end