rand(1:n, k) inside a Zygote gradient errors since 1.74.3 #854

Activity

1. CarloLucibello commented on Oct 24, 2026

   

   CarloLucibello

   on Oct 24, 2026

   MemberAuthor

   More actions

   The safest thing to do is to revert #604, and then to add a proper rule for 3-arg rand!, covering at least abstractarray 3rd argument and possibly generic collections. cc @devmotion
2. devmotion commented on Oct 24, 2026

   

   devmotion

   on Oct 24, 2026

   Member

   More actions

   Hmmm that's unfortunate. I made sure that all Zygote tests pass before merging #604. The previous situation with the too wide non-differentiable rules was worse though, IMO errors are much better than silently incorrect derivatives. It even seems that adding a non-differentiable rule for range sampling would also lead to silently incorrect derivatives as soon as you take derivatives with respect to a parameter of the range?
3. devmotion commented on Oct 24, 2026

   

   devmotion

   on Oct 24, 2026

   Member

   More actions

   Here are some concrete examples of why I think we shouldn't revert #604, nor add @non_differentiable rand!(::AbstractRNG, ::AbstractArray, ::AbstractRange).

   Sampling from a collection just indexes it at random indices. The indices depend only on the RNG state and the length of the collection, not on its values. So for a fixed seed, the derivative of the samples with respect to the collection's values is well defined and generally nonzero. The rule in ChainRules ≤ 1.74.2 (@non_differentiable rand!(::Any, ::Any, ::Any)) returned nothing for these cases without any error:

   using Zygote, Random, Distributions, LinearAlgebra  # ChainRules 1.74.2, Zygote 0.7.14, Julia 1.13.1
   
   Zygote.gradient(x -> sum(rand(Xoshiro(1), [x], 3)), 1.0)                   # (nothing,), correct: (3.0,)
   Zygote.gradient(x -> sum(rand(Xoshiro(1), LinRange(x, x + 1, 2), 3)), 1.0) # (nothing,), correct: (3.0,)
   Zygote.gradient(x -> sum(rand(Xoshiro(1), MvNormal(zeros(2), x * I))), 1.0) # (nothing,), correct: (0.2304...,)

   In the first two examples every sample is x or x + 1, so the correct gradient is 3 for any seed. The third is the example from #604 (comment); its correct value comes from central finite differences with the same seed. On 1.74.3 all three throw Mutating arrays is not supported, which is the better outcome.

   The proposed @non_differentiable rand!(::AbstractRNG, ::AbstractArray, ::AbstractRange) fixes rand(1:4, 5), but on 1.74.3 it also brings back the silently wrong result for the second example: (nothing,) instead of (3.0,). LinRange, StepRangeLen etc. have differentiable values.

   There's also a remaining rule that is still too wide on 1.74.3. Scalar sampling goes rand(rng, X) → rand(rng, Random.Sampler(rng, X, Val(1))), which hits @non_differentiable rand(::AbstractRNG, ::Random.Sampler) for every collection:

   # ChainRules 1.74.3
   Zygote.gradient(x -> rand(Xoshiro(1), [x]), 1.0)                   # (nothing,), correct: (1.0,)
   Zygote.gradient(x -> rand(Xoshiro(1), (x,)), 1.0)                  # (nothing,), correct: (1.0,)
   Zygote.gradient(x -> rand(Xoshiro(1), LinRange(x, x + 1, 2)), 1.0) # (nothing,), correct: (1.0,)

   My proposal:

   1. Restrict rand(::AbstractRNG, ::Random.Sampler) to the samplers that don't hold any differentiable data (SamplerType, SamplerTrivial{<:Random.FloatInterval}, integer range samplers, ...), so that the examples above error instead of returning nothing.
   2. Then add proper rules for sampling from arrays, ranges and tuples, rand(rng, X) and rand(rng, X, dims). The pullback scatters the cotangent back to the sampled indices, i.e. it reuses the getindex pullback. This fixes rand(1:4, 5) from this issue (and the example in the OP) without silently dropping gradients for LinRange & co.
4. devmotion

   closed this as completedin #855on Oct 25, 2026