Fix Zygote sol[:, j] (getindex ::Colon, ::Integer) returning zero gradient#1415
Merged
ChrisRackauckas merged 1 commit intoJul 2, 2026
Conversation
…dient Under RecursiveArrayTools v4, `AbstractVectorOfArray` getindex adjoints are inherited from Zygote's generic AbstractArray rules, so `sol[:, j]` (whole state at timestep j) had no VectorOfArray-specific adjoint and fell through to ChainRules' generic `∇getindex`, whose cotangent the solve reverse pass reads as zero. `Zygote.gradient` of e.g. `sum(sol[:, end])` therefore came back all zeros (silently wrong), while `sol[i]`, `sol[:, :]` and `sol.u[end]` were fine. Add an `@adjoint getindex(::ODESolution, ::Colon, ::Integer)` that scatters the cotangent into the j-th slot of a per-timestep cotangent, mirroring the existing `sol[i::Integer]` rule. The `::Colon` method is more specific than the existing `(sym, j::Integer)` method, so `sol[:, j]` dispatches to it. Verified (Julia 1.10.11): Zygote.gradient of `sum(solve(prob;p)[:, end])` and `[:, 2]` now match ForwardDiff (were `[0,0,0,0]`). Fixes the SciMLSensitivity Core1 "VecOfArray - Zygote" failure (SciMLSensitivity#1522). Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.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.
Ignore until reviewed by @ChrisRackauckas.
Under RecursiveArrayTools v4,
AbstractVectorOfArraygetindex adjoints are inherited from Zygote's genericAbstractArrayrules, sosol[:, j](whole state at timestepj) had no VectorOfArray-specific adjoint and fell through to ChainRules' generic∇getindex— whose cotangent the solve reverse pass reads as zero. SoZygote.gradientof e.g.sum(sol[:, end])came back all zeros (silently wrong), whilesol[i],sol[:, :], andsol.u[end]were fine.Adds an
@adjoint getindex(::ODESolution, ::Colon, ::Integer)that scatters the cotangent into thej-th slot of a per-timestep cotangent, mirroring the existingsol[i::Integer]rule.::Colonis more specific than the existing(sym, j::Integer)method, sosol[:, j]dispatches here.Verified (Julia 1.10.11):
Zygote.gradientofsum(solve(prob; p)[:, end])and[:, 2]now matchForwardDiff.gradient(were[0,0,0,0]);sol[:, :]/sol.u[end]unchanged. Fixes SciMLSensitivity #1522 (the Core1 "VecOfArray - Zygote" failure) — same family as #1410.🤖 Generated with Claude Code