From caf7d1d90acd26ed90a1be75be1c6bcce33cb369 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Sun, 25 Oct 2026 00:02:45 +0200 Subject: [PATCH 1/2] Restrict non-differentiable rand rules to data-free and integer samplers `@non_differentiable rand(::AbstractRNG, ::Random.Sampler)` matched every sampler, including the ones that sample from arrays, tuples, non-integer ranges and dictionaries. Zygote therefore silently returned `nothing` for e.g. `gradient(x -> rand([x]), 1.0)` instead of 1.0. The rule now only covers samplers whose output does not depend on any differentiable data (types, float intervals, BigFloat, Base strings and Char arrays) and samplers for ranges of Base integer types, for both `rand` and `rand!`. The latter fixes #854 without reintroducing the over-wide `rand!` rules removed in #604. The forwarding rules for `rand`/`rand!` and `Random.randexp(::AbstractRNG)` are removed: they are reached via the sampler rules or subsumed by `randexp(::Any...)`. `randstring` is extended to all argument types since its output is always a `String`. Co-Authored-By: Claude Opus 5.5 --- Project.toml | 2 +- src/rulesets/Random/random.jl | 30 ++++++++------- test/rulesets/Random/random.jl | 69 ++++++++++++++++++++++------------ 3 files changed, 63 insertions(+), 38 deletions(-) diff --git a/Project.toml b/Project.toml index aeb76428a..9cf6ab670 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "ChainRules" uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2" -version = "1.74.3" +version = "1.74.4" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" diff --git a/src/rulesets/Random/random.jl b/src/rulesets/Random/random.jl index 9a08f8551..ef9f67b0d 100644 --- a/src/rulesets/Random/random.jl +++ b/src/rulesets/Random/random.jl @@ -9,19 +9,23 @@ end @non_differentiable Broadcast.broadcastable(::AbstractRNG) -@non_differentiable Random.randexp(::AbstractRNG) -@non_differentiable Random.randstring(::AbstractRNG) - -@non_differentiable rand(::AbstractRNG, ::Random.Sampler) -@non_differentiable rand(::AbstractRNG, ::Integer...) -@non_differentiable rand(::Type{<:Real}) -@non_differentiable rand(::Type{<:Real}, ::Tuple) -@non_differentiable rand(::Type{<:Real}, ::Integer...) -@non_differentiable rand(::Integer...) - -@non_differentiable rand!(::AbstractArray) -@non_differentiable rand!(::AbstractRNG, ::AbstractArray) -@non_differentiable rand!(::AbstractRNG, ::AbstractArray, ::Type) +@non_differentiable Random.randstring(::Any...) + +# Only samplers whose output does not depend on any differentiable data, or that sample from +# ranges of Base integer types: e.g. not the ones for sampling from arrays, non-integer ranges, +# tuples, or user-defined types. +const NonDifferentiableSampler = Union{ + Random.SamplerType, + Random.SamplerTrivial{<:Random.FloatInterval}, + Random.SamplerBigFloat, + Random.SamplerSimple{<:Union{String,SubString{String},AbstractArray{Char}}}, + Random.SamplerRangeInt{<:Base.BitInteger}, + Random.SamplerRangeFast, + Random.SamplerRangeNDL{<:Any,<:Base.BitInteger}, + Random.SamplerBigInt, +} +@non_differentiable rand(::AbstractRNG, ::NonDifferentiableSampler) +@non_differentiable rand!(::AbstractRNG, ::AbstractArray, ::NonDifferentiableSampler) @non_differentiable randexp(::Any...) diff --git a/test/rulesets/Random/random.jl b/test/rulesets/Random/random.jl index 6915a3f86..1f2360898 100644 --- a/test/rulesets/Random/random.jl +++ b/test/rulesets/Random/random.jl @@ -32,30 +32,6 @@ Random.rand(d::NormalDistribution) = d.μ + d.σ*randn() end @testset "rand" begin - non_differentiables = [ - ((), Float64), - ((MersenneTwister(123),), Float64), - ((MersenneTwister(123),2,2), Matrix{<:Float64}), - ((Float32,), Float32), - ((Float32,2,2), Matrix{<:Float32}), - ((Float32,(2,2)), Matrix{<:Float32}), - ((2,2), Matrix{<:Float64}), - ] - push!(non_differentiables, ((Xoshiro(123),), Float64)) - push!(non_differentiables, ((Xoshiro(123), 2, 2), Matrix{<:Float64})) - - for (args, xType) in non_differentiables - x, dΩ = frule((ZeroTangent(), randn(args...)), rand, args...) - @test x isa xType - @test dΩ isa NoTangent - - x, pb = rrule(rand, args...) - @test x isa xType - dself, dargs = Iterators.peel(pb(10.0)) - @test iszero(dself) - @test all(darg isa NoTangent for darg in dargs) - end - # Make sure that we do *not* have these set as non_differentiable. as they are differentiable @test nothing === frule( (ZeroTangent(), Tangent{NormalDistribution}(μ=0.5,σ=2.0)), @@ -64,4 +40,49 @@ Random.rand(d::NormalDistribution) = d.μ + d.σ*randn() ) @test rrule(rand, NormalDistribution(0.1,1.5)) === nothing end + + @testset "randstring" begin + for args in ((), (5,), (Xoshiro(123),), (Xoshiro(123), 5), ("ab",), (Xoshiro(123), 'a':'z', 5)) + s, ds = frule((NoTangent(), map(Returns(NoTangent()), args)...), randstring, args...) + @test s isa String + @test ds isa NoTangent + s, pb = rrule(randstring, args...) + @test s isa String + @test all(darg isa NoTangent for darg in pb(NoTangent())) + end + end + + @testset "rand(rng, ::Sampler)" begin + rng = Xoshiro(123) + + # Samplers without any differentiable data and samplers for integer ranges are non-differentiable + for X in ( + Float64, Float32, Int, Bool, ComplexF64, BigFloat, "abc", SubString("abc", 2), 'a':'z', ['a', 'b'], + 1:4, Base.OneTo(4), Int8(1):Int8(4), UInt128(1):UInt128(4), typemin(Int):typemax(Int), big(1):big(4), + ) + sp = Random.Sampler(rng, X, Val(1)) + x, dx = frule((NoTangent(), NoTangent(), NoTangent()), rand, rng, sp) + @test x isa eltype(sp) + @test dx isa NoTangent + x, pb = rrule(rand, rng, sp) + @test x isa eltype(sp) + @test all(darg isa NoTangent for darg in pb(1.0)) + + A = Vector{eltype(sp)}(undef, 3) + spInf = Random.Sampler(rng, X, Val(Inf)) + Ω, dΩ = frule((NoTangent(), NoTangent(), NoTangent(), NoTangent()), rand!, rng, A, spInf) + @test Ω === A + @test dΩ isa NoTangent + Ω, pb! = rrule(rand!, rng, A, spInf) + @test Ω === A + @test all(darg isa NoTangent for darg in pb!(ones(3))) + end + + # Samplers for collections with possibly differentiable elements are not + for X in ([1.0, 2.0], (1.0,), (1.0, 2.0), LinRange(0.0, 1.0, 3), 0.0:0.5:1.0, UnitRange(1.0, 4.0), Dict(1 => 2.0), Set([1.0])) + @test rrule(rand, rng, Random.Sampler(rng, X, Val(1))) === nothing + @test rrule(rand!, rng, Vector{Any}(undef, 3), Random.Sampler(rng, X, Val(Inf))) === nothing + end + @test rrule(rand, rng, Random.SamplerRangeNDL(1.5, UInt64(4))) === nothing + end end From 851e689e12cc04fb860893198d1c0311ddb0bf8a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Sun, 25 Oct 2026 00:27:08 +0200 Subject: [PATCH 2/2] Apply batched suggestions from code review Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- test/rulesets/Random/random.jl | 50 +++++++++++++++++++++++++++++----- 1 file changed, 43 insertions(+), 7 deletions(-) diff --git a/test/rulesets/Random/random.jl b/test/rulesets/Random/random.jl index 1f2360898..dbff1c824 100644 --- a/test/rulesets/Random/random.jl +++ b/test/rulesets/Random/random.jl @@ -42,8 +42,17 @@ Random.rand(d::NormalDistribution) = d.μ + d.σ*randn() end @testset "randstring" begin - for args in ((), (5,), (Xoshiro(123),), (Xoshiro(123), 5), ("ab",), (Xoshiro(123), 'a':'z', 5)) - s, ds = frule((NoTangent(), map(Returns(NoTangent()), args)...), randstring, args...) + for args in ( + (), + (5,), + (Xoshiro(123),), + (Xoshiro(123), 5), + ("ab",), + (Xoshiro(123), 'a':'z', 5), + ) + s, ds = frule( + (NoTangent(), map(Returns(NoTangent()), args)...), randstring, args... + ) @test s isa String @test ds isa NoTangent s, pb = rrule(randstring, args...) @@ -57,8 +66,22 @@ Random.rand(d::NormalDistribution) = d.μ + d.σ*randn() # Samplers without any differentiable data and samplers for integer ranges are non-differentiable for X in ( - Float64, Float32, Int, Bool, ComplexF64, BigFloat, "abc", SubString("abc", 2), 'a':'z', ['a', 'b'], - 1:4, Base.OneTo(4), Int8(1):Int8(4), UInt128(1):UInt128(4), typemin(Int):typemax(Int), big(1):big(4), + Float64, + Float32, + Int, + Bool, + ComplexF64, + BigFloat, + "abc", + SubString("abc", 2), + 'a':'z', + ['a', 'b'], + 1:4, + Base.OneTo(4), + Int8(1):Int8(4), + UInt128(1):UInt128(4), + typemin(Int):typemax(Int), + big(1):big(4), ) sp = Random.Sampler(rng, X, Val(1)) x, dx = frule((NoTangent(), NoTangent(), NoTangent()), rand, rng, sp) @@ -70,7 +93,9 @@ Random.rand(d::NormalDistribution) = d.μ + d.σ*randn() A = Vector{eltype(sp)}(undef, 3) spInf = Random.Sampler(rng, X, Val(Inf)) - Ω, dΩ = frule((NoTangent(), NoTangent(), NoTangent(), NoTangent()), rand!, rng, A, spInf) + Ω, dΩ = frule( + (NoTangent(), NoTangent(), NoTangent(), NoTangent()), rand!, rng, A, spInf + ) @test Ω === A @test dΩ isa NoTangent Ω, pb! = rrule(rand!, rng, A, spInf) @@ -79,9 +104,20 @@ Random.rand(d::NormalDistribution) = d.μ + d.σ*randn() end # Samplers for collections with possibly differentiable elements are not - for X in ([1.0, 2.0], (1.0,), (1.0, 2.0), LinRange(0.0, 1.0, 3), 0.0:0.5:1.0, UnitRange(1.0, 4.0), Dict(1 => 2.0), Set([1.0])) + for X in ( + [1.0, 2.0], + (1.0,), + (1.0, 2.0), + LinRange(0.0, 1.0, 3), + 0.0:0.5:1.0, + UnitRange(1.0, 4.0), + Dict(1 => 2.0), + Set([1.0]), + ) @test rrule(rand, rng, Random.Sampler(rng, X, Val(1))) === nothing - @test rrule(rand!, rng, Vector{Any}(undef, 3), Random.Sampler(rng, X, Val(Inf))) === nothing + @test rrule( + rand!, rng, Vector{Any}(undef, 3), Random.Sampler(rng, X, Val(Inf)) + ) === nothing end @test rrule(rand, rng, Random.SamplerRangeNDL(1.5, UInt64(4))) === nothing end