fix: incorrect scope for seed (#1343)

* fix: incorrect scope

* fix: use copyto!

Author: avik-pal

src/stdlibs/Random.jl

@@ -43,7 +43,7 @@ end
 end
 
 @noinline function Random.seed!(rng::TracedRNG, seed::TracedRArray{UInt64,1})
-    rng.seed.mlir_data = seed.mlir_data
+    copyto!(rng.seed, seed)
     return rng
 end
 
@@ -82,7 +82,7 @@ Base.copy(rng::TracedRNG) = TracedRNG(copy(rng.seed), rng.algorithm)
 ) where {T,N}
     length(A) == 0 && return A
     res = Ops.rng_bit_generator(T, rng.seed, [size(A)...]; rng.algorithm)
-    rng.seed.mlir_data = res.output_state.mlir_data
+    copyto!(rng.seed, res.output_state)
     TracedUtils.set_mlir_data!(A, res.output.mlir_data)
     return A
 end
@@ -92,7 +92,7 @@ end
 ) where {T,N}
     length(A) == 0 && return A
     res = Ops.randn(T, rng.seed, [size(A)...]; rng.algorithm)
-    rng.seed.mlir_data = res.output_state.mlir_data
+    copyto!(rng.seed, res.output_state)
     TracedUtils.set_mlir_data!(A, res.output.mlir_data)
     return A
 end
@@ -102,7 +102,7 @@ end
 ) where {T,N}
     length(A) == 0 && return A
     res = Ops.randexp(T, rng.seed, [size(A)...]; rng.algorithm)
-    rng.seed.mlir_data = res.output_state.mlir_data
+    copyto!(rng.seed, res.output_state)
     TracedUtils.set_mlir_data!(A, res.output.mlir_data)
     return A
 end

test/autodiff.jl

@@ -1,4 +1,4 @@
-using Enzyme, Reactant, Test
+using Enzyme, Reactant, Test, Random
 
 square(x) = x * 2
 
@@ -228,3 +228,30 @@ vector_forward_ad(x) = Enzyme.autodiff(Forward, fn, BatchDuplicated(x, Enzyme.on
     @test res[1][3] ≈ res_enz[1][3]
     @test res[1][4] ≈ res_enz[1][4]
 end
+
+function simple_forward(x, st)
+    rng = copy(st.rng)
+    y = similar(x)
+    rand!(rng, y)
+    return x .+ y, (; rng)
+end
+
+function gradient_fn(x, st)
+    stₙ = Ref{Any}(nothing)
+    function lfn(x, st_old)
+        y, st_new = simple_forward(x, st_old)
+        stₙ[] = st_new
+        return sum(abs2, y)
+    end
+    return Enzyme.gradient(Reverse, lfn, x, Const(st)), stₙ[]
+end
+
+@testset "seed" begin
+    x = Reactant.to_rarray(rand(2, 2))
+    st = (; rng=Reactant.ConcreteRNG())
+
+    @test begin
+        hlo = @code_hlo gradient_fn(x, st)
+        contains(repr(hlo), "stablehlo.rng_bit_generator")
+    end
+end