NonlinearSolveBase: bring solve_up Enzyme rule to parity with DiffEqBase (MTK DAE init under runtime activity) - #944
Merged
ChrisRackauckas merged 1 commit intoMay 30, 2026
Conversation
NonlinearSolveBase's `solve_up` Enzyme rule was a clone of the older DiffEqBase
`solve_up` rule and is missing the same modernizations DiffEqBase has since
received. Reverse-mode Enzyme through an MTK DAE solve (whose initialization
solves a `NonlinearProblem`) under `set_runtime_activity` hits, in sequence:
1. `NonConstantKeywordArgException` — no `inactive_kwarg` declaration. `solve_up`'s
differentiable inputs (prob/u0/p) are positional and its kwargs are solver
config, so declare them inactive (mirrors DiffEqBase).
2. `MethodError` on `augmented_primal`/`reverse` — the rule required
`sensealg::Const`, but under runtime activity a `SteadyStateAdjoint` sensealg
arrives `Duplicated`. Widen to `Enzyme.Annotation{<:Union{Nothing,
AbstractSensitivityAlgorithm}}` and skip the sensealg slot in reverse
(mirrors SciML/OrdinaryDiffEq#3678/#3681).
3. `MethodError: MixedDuplicated(::NonlinearSolution)` — the shadow must be
`Ref`-wrapped for the abstract MixedDuplicated case; dereference it in reverse
(mirrors SciML/OrdinaryDiffEq#3700).
4. `ReverseRuleReturnError` — the init polyalgorithm is promoted to `Active`
under runtime activity, so the reverse rule must return a cotangent
(`make_zero`) for `Active` args rather than `nothing`.
Each change was confirmed to advance past its error in sequence. The fixes are
no-ops for the existing concrete-`Const`-sensealg / non-abstract paths.
Companion to SciML/SciMLSensitivity#1463 (the SciMLSensitivity-side fixes for the
same MTK DAE + GaussAdjoint(EnzymeVJP()) chain).
Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
ChrisRackauckas
marked this pull request as ready for review
May 30, 2026 21:35
ChrisRackauckas
added a commit
that referenced
this pull request
May 31, 2026
…nator adjoint path (#946) Follow-up to #944. When Enzyme reverse-mode differentiates through an MTK DAE solve, the initialization solves a `NonlinearProblem` via `solve_up`. The custom `solve_up` Enzyme rule re-runs the solve through `_solve_adjoint` -> `get_concrete_problem`, which `maybe_wrap_nonlinear_f`-wraps the IIP function in a `FunctionWrappersWrapper` (`AutoSpecializeCallable`). That wrapping is type-unstable: Enzyme's traced forward solve unwraps it (`maybe_unwrap_prob_for_enzyme`, #940) so its inferred return type carries the bare function, but the adjoint re-wraps it, so the rule's returned primal type no longer matches and Enzyme aborts with `EnzymeRuntimeException: Expected return type of primal to be NonlinearSolution{...}`. `maybe_unwrap_prob_for_enzyme` keys off the solver's own autodiff, which for an MTK DAE init is ForwardDiff even when the outer differentiation is Enzyme, so it does not fire on this path. Key off the originator instead: unwrap `_prob.f.f` via `get_raw_f` when `originator isa EnzymeOriginator`, matching the (unwrapped) type Enzyme's traced forward produces. Companion to SciML/SciMLSensitivity#1463. Co-authored-by: ChrisRackauckas-Claude <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.
Please ignore until reviewed by @ChrisRackauckas.
Summary
NonlinearSolveBase'ssolve_upEnzyme rule (NonlinearSolveBaseEnzymeExt.jl) is a clone of the older DiffEqBasesolve_uprule and is missing the modernizations DiffEqBase has since received. This surfaces when reverse-mode Enzyme runs through an MTK DAE solve whose initialization solves aNonlinearProblem(e.g.BrownFullBasicInit/DefaultInit) underset_runtime_activity— exactly theSciMLSensitivitytest/mtk.jlEnzyme path. The rule hits four errors in sequence; each fix advances past the previous error:NonConstantKeywordArgException— noinactive_kwargdeclaration.solve_up's differentiable inputs (prob/u0/p) are positional; kwargs are solver config → declare them inactive (mirrors the DiffEqBase declaration).MethodErroronaugmented_primal/reverse— the rule requiredsensealg::Const, but under runtime activity aSteadyStateAdjointsensealg arrivesDuplicated. Widen toEnzyme.Annotation{<:Union{Nothing, AbstractSensitivityAlgorithm}}and skip the sensealg slot in reverse (mirrors OrdinaryDiffEq #3678/#3681).MethodError: MixedDuplicated(::NonlinearSolution)— the shadow must beRef-wrapped for the abstract-MixedDuplicated case; dereference in reverse (mirrors OrdinaryDiffEq #3700).ReverseRuleReturnError— the auto-selected init polyalgorithm is promoted toActiveunder runtime activity, so the reverse rule must return a cotangent (make_zero) forActiveargs, notnothing.The changes are no-ops for the existing concrete-
Const-sensealg / non-abstract paths.Verification
solve_upfixes that are gradient-verified (OrdinaryDiffEq #3700 / SciMLSensitivity #1463 reproduce a bit-for-bit / ForwardDiff-matching gradient).BrownFullBasicInitDAE-init gradient check through this path (vs ForwardDiff) is running; I'll post the result.Relationship / remaining blocker
Companion to SciML/SciMLSensitivity#1463 (the SciMLSensitivity-side fixes for the same MTK DAE +
GaussAdjoint(EnzymeVJP())chain).SciMLSensitivitytest/mtk.jlstill does not fully pass after these fixes, but for a separate, backend-independent reason: several of its setups useguesses = [w2 => 0.0]with the constraint0 = x² + y² - w2², making the algebraic-constraint Jacobiandhda = -2·w2singular in the GaussAdjoint DAE adjoint (SciMLSensitivity/src/adjoint_common.jl:770,dhda' \ gᵤ). ThatSingularExceptionis hit by Zygote/ReverseDiff through the sameGaussAdjointtoo (verification in progress) — i.e. it is not related to these Enzyme-rule fixes.🤖 Generated with Claude Code