fix: fixes for random sampling (#1513)

Author: avik-pal

src/Overlay.jl

@@ -43,8 +43,8 @@ for randfun in (:rand, :randn, :randexp)
         @reactant_overlay @noinline function Random.$(randfun)(
             rng::AbstractRNG, ::Type{T}, dims::Dims
         ) where {T}
-            if T <: ReactantPrimitive
-                return TracedRandom.$(overload_randfun)(rng, T, dims)
+            if unwrapped_eltype(T) <: ReactantPrimitive
+                return TracedRandom.$(overload_randfun)(rng, unwrapped_eltype(T), dims)
             end
             @warn "Reactant doesn't support sampling of $(T) with the current \
                    interpreter. Falling back to native interpreter." maxlog = 1
@@ -60,8 +60,10 @@ for randfun in (:rand, :randn, :randexp)
         @reactant_overlay @noinline function Random.$(randfun)(
             rng::AbstractRNG, ::Type{T}, dim1::Integer, dims::Integer...
         ) where {T}
-            if T <: ReactantPrimitive
-                return TracedRandom.$(overload_randfun)(rng, T, dim1, dims...)
+            if unwrapped_eltype(T) <: ReactantPrimitive
+                return TracedRandom.$(overload_randfun)(
+                    rng, unwrapped_eltype(T), dim1, dims...
+                )
             end
             @warn "Reactant doesn't support sampling of $(T) with the current \
                    interpreter. Falling back to native interpreter." maxlog = 1
@@ -72,8 +74,8 @@ for randfun in (:rand, :randn, :randexp)
         @reactant_overlay @noinline function Random.$(randfun)(
             rng::AbstractRNG, ::Type{T}=Float64
         ) where {T}
-            if T <: ReactantPrimitive
-                return TracedRandom.$(overload_randfun)(rng, T)
+            if unwrapped_eltype(T) <: ReactantPrimitive
+                return TracedRandom.$(overload_randfun)(rng, unwrapped_eltype(T))
             end
             @warn "Reactant doesn't support sampling of $(T) with the current \
                    interpreter. Falling back to native interpreter." maxlog = 1
@@ -86,6 +88,9 @@ for randfun in (:rand, :randn, :randexp)
         )
             return TracedRandom.$(overload_randfun!)(rng, A)
         end
+        @reactant_overlay @noinline function Random.$(randfun!)(A::AnyTracedRArray)
+            return TracedRandom.$(overload_randfun!)(TracedRandom.default_rng(), A)
+        end
     end
 end
 

test/integration/random.jl

@@ -185,3 +185,19 @@ end
         @test μ ≈ 1.0 atol = 0.05 rtol = 0.05
     end
 end
+
+rand_sample(rng, x) = rand(rng, eltype(x), size(x))
+
+function rand_on_device()
+    x = Reactant.Ops.fill(0.0f0, (3, 4, 5))
+    rand!(x)
+    return x
+end
+
+@testset "TracedTypes in Sampling" begin
+    @test @jit(rand_sample(Reactant.ReactantRNG(), rand(3, 4))) isa
+        ConcreteRArray{Float64,2}
+    @test @jit(rand_sample(Reactant.ReactantRNG(), Reactant.to_rarray(rand(3, 4)))) isa
+        ConcreteRArray{Float64,2}
+    @test @jit(rand_on_device()) isa ConcreteRArray{Float32,3}
+end