Fix Mooncake gradient crash with ComponentArray#1565
Open
AstitvaAggarwal wants to merge 2 commits into
SciML:masterSciML/SciMLSensitivity.jl:masterfrom
AstitvaAggarwal:fix/1548-componentarray-mooncake-gradientAstitvaAggarwal/SciMLSensitivity.jl:fix/1548-componentarray-mooncake-gradientCopy head branch name to clipboard
Open
Fix Mooncake gradient crash with ComponentArray#1565AstitvaAggarwal wants to merge 2 commits intoSciML:masterSciML/SciMLSensitivity.jl:masterfrom AstitvaAggarwal:fix/1548-componentarray-mooncake-gradientAstitvaAggarwal/SciMLSensitivity.jl:fix/1548-componentarray-mooncake-gradientCopy head branch name to clipboard
AstitvaAggarwal wants to merge 2 commits into
SciML:masterSciML/SciMLSensitivity.jl:masterfrom
AstitvaAggarwal:fix/1548-componentarray-mooncake-gradientAstitvaAggarwal/SciMLSensitivity.jl:fix/1548-componentarray-mooncake-gradientCopy head branch name to clipboard
Conversation
`concrete_solve.jl` assumed a wrapped-array tangent always exposes its data via a `.x` field, which only held for the package's own internal test wrapper type. For `ComponentArrays.ComponentVector`, that grabbed the wrong backing field (a `ZeroTangent`), causing `MethodError: no method matching vec(::ChainRulesCore.ZeroTangent)`. Add `_tangent_array_data` to find the tangent's one non-zero field instead of assuming a fixed name. Fixes SciML#1548
AstitvaAggarwal
marked this pull request as draft
July 22, 2026 00:51
Tangent overloads getproperty to mirror the primal's fields, so `v.backing` doesn't return the Tangent's own backing field - it returns ZeroTangent() (the overload's fallback for a missing property). Use getfield(v, :backing) instead, which bypasses the overload safely and doesn't require importing the non-public ChainRulesCore.backing function (which was failing Aqua's ExplicitImports public-API check in CI).
AstitvaAggarwal
marked this pull request as ready for review
July 22, 2026 12:28
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
1 participant
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.
Summary
concrete_solve.jl'sForwardDiffSensitivity/ForwardSensitivityadjoint pullback assumed a wrapped-array tangent always exposes its data via a.xfield. That assumption only held for the package's own internal test wrapper type — forComponentArrays.ComponentVector, it grabbed the wrong backing field (aZeroTangent), causing:when computing a Mooncake gradient of an ODE solve with a
ComponentVectoru0(see #1548 for the full MRE/stacktrace).This is a generic
ChainRulesCore.Tangent-consuming code path shared by every AD backend that calls into it viaChainRulesCore.rrule(Mooncake included, via its genericrrule-wrapping bridge). Empirically, Zygote/ReverseDiff never construct aTangentwith this exact field-mirroring shape for this scenario, so Mooncake is the only backend that currently trips it, but the fix itself is backend-agnostic.Added
_tangent_array_data(v::Tangent), which finds the tangent's one non-zero backing field instead of assuming a fixed name (recursing if that field is itself aTangent), and used it at both call sites that previously didv.x.Test plan
test/Core1/mooncake_componentarray_gradient.jl, mirroring the existingmooncake_vjp_prob_kwargs.jlstyle, registered intest/runtests.jlCloses #1548
Cc: @ChrisRackauckas