test_rrule: also check output tangents with thunked fields #296

check_thunked_output_tangent=true only checks the pullback with the whole output tangent thunked, i.e. pullback(@thunk(ȳ)). But AD systems also pass structural tangents whose fields are thunks. For example, Zygote's tuple getindex adjoint keeps thunks, so for f(x)[i] the rule receives Tangent{Any}((ZeroTangent(), …, thunk, …)).

Rules that handle the outer thunk correctly can still fail on this. In JuliaDiff/ChainRules.jl#843, the findnz rule passed test_rrule, but with Zygote it failed:

using SparseArrays, Zygote
Zygote.gradient(A -> sum(findnz(A)[3]), sprand(5, 5, 0.4))
# MethodError: no method matching sparse(::Vector{Int64}, ::Vector{Int64}, ::InplaceableThunk{...}, ::Int64, ::Int64)

The pullback destructured _, _, V̄ = unthunk(Δ). Destructuring a Tangent does not unthunk its fields, unlike getindex/getproperty (see JuliaDiff/ChainRulesCore.jl#712).

Proposal: when the output tangent ȳ is a Tangent, or a Tuple/NamedTuple that is converted to one, test_rrule should additionally call the pullback with every field wrapped in a thunk, e.g. Tangent{P}(map(x -> @thunk(x), backing(ȳ))...), and compare the result with the unthunked cotangents. This could reuse the existing check_thunked_output_tangent flag or get a new one.