[devmotion]

Fixes #854 without reverting #604, and removes the remaining silently-incorrect rand rule.

Problem

@non_differentiable rand(::AbstractRNG, ::Random.Sampler) matches every sampler, including the ones that sample from collections of differentiable values (SamplerSimple{<:Array}, SamplerTrivial{<:Tuple}, SamplerSimple{<:LinRange}, SamplerSimple{<:Dict}, SamplerTag{<:Set}, ...). Hence, e.g., on ChainRules 1.74.3

julia> Zygote.gradient(x -> rand([x]), 1.0)
(nothing,)  # correct: (1.0,)

At the same time, #604 removed the over-wide rand! rules (which caused the same problem for e.g. sum(rand([x], 3))), which made rand(1:4, 5) error in Zygote (#854).

Changes

• rand(::AbstractRNG, sp) and rand!(::AbstractRNG, ::AbstractArray, sp) are only non-differentiable for samplers whose output does not depend on any differentiable data or that sample from ranges of Base integer types:
  • SamplerType, SamplerTrivial{<:FloatInterval}, SamplerBigFloat
  • SamplerSimple for String, SubString{String} and AbstractArray{Char}
  • SamplerRangeInt{<:BitInteger}, SamplerRangeFast, SamplerRangeNDL{<:Any,<:BitInteger}, SamplerBigInt
• The rand! rule fixes rand(1:n, k) inside a Zygote gradient errors since 1.74.3 #854 (rand(1:4, 5), rand!(rng, A, 1:4)).
• Removed redundant rules:
  • the forwarding rules rand(::Integer...), rand(::Type{<:Real}, ...), rand(::AbstractRNG, ::Integer...), rand!(::AbstractArray), rand!(::AbstractRNG, ::AbstractArray), rand!(::AbstractRNG, ::AbstractArray, ::Type): these calls reach the sampler rules above.
  • Random.randexp(::AbstractRNG): subsumed by randexp(::Any...).
• randstring(::AbstractRNG) → randstring(::Any...) (output is always a String); randstring() and randstring(rng, n) errored in Zygote before.

Zygote results (Julia 1.10.12 and 1.13.1)

	1.74.3	this PR
rand(rng, [x]), rand(rng, (x,)), rand(rng, (x, x))	nothing ❌	1.0 ✅
rand(rng, LinRange(x, x + 1, 2)), Dict, Set	nothing ❌	error
sum(x .* rand(rng, 1:4, 5)) (#854)	error	13.0 ✅
x * rand(rng, 1:4), UInt128 ranges	✅	✅
rand(rng), rand(rng, 3), rand!(rng, zeros(3)), rand(Float32, 2, 2), rand(Int), BigFloat, 'a':'z', randstring(rng), ...	✅	✅

Proper (gather/scatter) rules for sampling from collections of differentiable values can be added in a follow-up PR.

Notes

• Integer-valued samples are treated as non-differentiable, consistent with how Mooncake and Enzyme treat integers.
• rrule/frule called directly on the removed forwarders (e.g. rrule(rand, 2, 2)) now return nothing; AD systems reach the sampler rules instead. I checked Zygote but not Diffractor.
• BigInt ranges (rand(rng, big(1):big(4))) still error in Zygote, as on 1.74.3: Zygote fails in the construction of the sampler, before the rand rule applies.

🤖 Generated with Claude Code

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 96.57%. Comparing base (e6e8fae) to head (851e689).

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #855      +/-   ##
==========================================
- Coverage   96.57%   96.57%   -0.01%     
==========================================
  Files          28       28              
  Lines        3618     3616       -2     
==========================================
- Hits         3494     3492       -2     
  Misses        124      124              

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[CarloLucibello]

This seems wide enough, thanks!