Restore non-differentiable rand!(::AbstractRNG, ::AbstractArray)

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRules"
 uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2"
-version = "1.74.11"
+version = "1.74.12"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/rulesets/Random/random.jl

@@ -26,6 +26,8 @@ const NonDifferentiableSampler = Union{
 }
 @non_differentiable rand(::AbstractRNG, ::NonDifferentiableSampler)
 @non_differentiable rand!(::AbstractRNG, ::AbstractArray, ::NonDifferentiableSampler)
+# GPUArrays and CUDA implement this directly, without a sampler
+@non_differentiable rand!(::AbstractRNG, ::AbstractArray)
 
 @non_differentiable randexp(::Any...)
 

test/rulesets/Random/random.jl

@@ -1,4 +1,4 @@
-using ChainRules, ChainRulesCore, Random, Test
+using ChainRules, ChainRulesCore, JLArrays, Random, Test
 
 # Simple Distributions like object for testing purposes
 struct NormalDistribution
@@ -7,6 +7,10 @@ struct NormalDistribution
 end
 Random.rand(d::NormalDistribution) = d.μ + d.σ*randn()
 
+# RNG that, like GPUArrays and CUDA, implements `rand!(rng, A)` without a sampler
+struct DirectRNG <: AbstractRNG end
+Random.rand!(::DirectRNG, A::Array) = fill!(A, 0.5)
+
 @testset "Random" begin
     rng_types = [MersenneTwister, Xoshiro]
 
@@ -121,4 +125,19 @@ Random.rand(d::NormalDistribution) = d.μ + d.σ*randn()
         end
         @test rrule(rand, rng, Random.SamplerRangeNDL(1.5, UInt64(4))) === nothing
     end
+
+    @testset "rand!(rng, A)" begin
+        for (rng, A) in (
+            (Xoshiro(123), zeros(3)),
+            (Xoshiro(123), JLArray(zeros(Float32, 3))),  # GPUArrays: no sampler
+            (DirectRNG(), zeros(3)),
+        )
+            Ω, dΩ = frule((NoTangent(), NoTangent(), NoTangent()), rand!, rng, A)
+            @test Ω === A
+            @test dΩ isa NoTangent
+            Ω, pb! = rrule(rand!, rng, A)
+            @test Ω === A
+            @test all(darg isa NoTangent for darg in pb!(ones(3)))
+        end
+    end
 end