fix: handle ZeroTangent/AbstractZero in RecursiveArrayToolsZygoteExt adjoint#606
Merged
ChrisRackauckas merged 1 commit intoJun 6, 2026
Conversation
…adjoint The Zygote pullbacks for VectorOfArray/DiffEqArray (vofa_u_adjoint) and ArrayPartition only special-cased `isnothing(d_i)` for missing per-slice gradients. A ChainRulesCore.ZeroTangent (AbstractZero) cotangent element was passed through unchanged, and the subsequent VectorOfArray/DiffEqArray construction called `size(::ZeroTangent)`, throwing a MethodError. This broke reverse-mode AD through DiffEqArray(vec, ts) (e.g. DifferenceEquations.jl test/linear_gradients.jl). Treat AbstractZero the same as nothing: a structural zero gradient for that slice. Added a regression test in test/adjoints.jl exercising a ZeroTangent cotangent element for both VectorOfArray and DiffEqArray. Regression test run locally on Julia 1.10: PASS. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
The Zygote pullbacks in
ext/RecursiveArrayToolsZygoteExt.jlforVectorOfArray/DiffEqArray(vofa_u_adjoint) andArrayPartitiononly special-casedisnothing(d_i)for missing per-slice gradients. When a cotangent element is aChainRulesCore.ZeroTangent(anAbstractZero, a structural zero gradient), it was passed through unchanged. The subsequentVectorOfArray/DiffEqArrayconstruction then callssize(::ZeroTangent), throwing:This breaks reverse-mode AD through
DiffEqArray(vec, ts)and shows up downstream (e.g. DifferenceEquations.jltest/linear_gradients.jl:67).Fix
Treat
AbstractZerothe same asnothingin the adjoints: a structural zero gradient for that slice (skip reconstruction, propagate a real zero array).AbstractZerois imported from the already-availableChainRulesCore.Test
Added a focused regression test in
test/adjoints.jlthat feeds aZeroTangentcotangent element throughvofa_u_adjointfor bothVectorOfArrayandDiffEqArrayand asserts the result is the correct zero gradient (and no longer errors). Verified the test fails on the unpatched code and passes with the fix.Local run on Julia 1.10 (
test/adjoints.jl): PASS.🤖 Generated with Claude Code