NonlinearSolveBase: mark nonlinear-solve algorithms Enzyme-inactive - #947
Merged
ChrisRackauckas merged 1 commit intoMay 31, 2026
Conversation
The default `NonlinearSolvePolyAlgorithm` reaches the `solve_up` custom Enzyme
rule as a positional argument when a `NonlinearProblem` is solved under
`Enzyme.set_runtime_activity(Reverse)` — e.g. the ModelingToolkit DAE
initialization solve reached through `get_initial_values`/`solve_up`. Nonlinear
algorithms are pure solver configuration (Vals, Rationals, Nothing/Missing
sentinels, nested immutable sub-algorithms, a ForwardDiff tag type parameter)
with no differentiable floating-point data, but without an inactivity
declaration Enzyme promotes that configuration argument to `Duplicated`. That
trips the `roots_activep != activep` assertion in `enzyme_custom_setup_args`,
or leaves a spurious algorithm shadow that corrupts activity bookkeeping for
the genuinely-active `u0`/`p` (silently dropping a cotangent component).
Declare the whole `AbstractNonlinearSolveAlgorithm` hierarchy `inactive_type`,
so Enzyme treats it as `Const` (`activep == roots_activep == DFT_CONSTANT`).
Mirrors SciMLBase's `inactive_type(::Type{<:AbstractSensitivityAlgorithm})` and
LinearSolve's `inactive_type(::Type{<:SciMLLinearSolveAlgorithm})`.
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 31, 2026 18:31
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
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.
What
Declare the whole
AbstractNonlinearSolveAlgorithmhierarchyEnzyme.EnzymeRules.inactive_type, inNonlinearSolveBaseEnzymeExt:Why
When a
NonlinearProblemis solved underEnzyme.set_runtime_activity(Reverse)— most notably the ModelingToolkit DAE-initialization solve reached throughget_initial_values→solve_up— the defaultNonlinearSolvePolyAlgorithmarrives at thesolve_upcustom Enzyme rule as a positionalargs.... Nonlinear-solve algorithms are pure solver configuration (Vals,Rationals,Nothing/Missingsentinels, nested immutable sub-algorithms, and aForwardDifftag type parameter) with no differentiable floating-point data, but without an inactivity declarationset_runtime_activitypromotes that configuration argument toDuplicated. That manifests, nondeterministically, as either:roots_activep (DFT_CONSTANT) != activep (DFT_DUP_ARG)assertion inside Enzyme'senzyme_custom_setup_args, oru0/p, silently dropping a cotangent component.Marking the algorithm hierarchy inactive forces
Const, soactivep == roots_activep == DFT_CONSTANT. This mirrors the established, shipping pattern inSciMLBase(inactive_type(::Type{<:AbstractSensitivityAlgorithm})) andLinearSolve(inactive_type(::Type{<:SciMLLinearSolveAlgorithm})).This is type-safe (the abstract supertype covers the polyalg wrapper, every nested member algorithm, and the bare-algorithm case) and semantically exact (algorithms are never a differentiation target — gradients flow only through
prob/u0/p). It is precompile-safe: a more-specific method on an already-loadedEnzymeRules.inactive_type, no parent-module method overwritten.Context
This is the NonlinearSolveBase half of getting Enzyme reverse-mode AD to work on the outside of an MTK DAE solve (the
within_autodiffwrap-skip and theinactive_kwargfor thesolve_uprule are already merged). A downstream SciMLSensitivity PR will flip the correspondingmtk.jltest on and bump its lower bound once this is tagged.Test
Adds
lib/NonlinearSolveBase/test/enzyme_inactive_algorithm.jl(wired intoruntests.jl), a deterministic regression test that the declaration coversAbstractNonlinearSolveAlgorithm, the concreteNonlinearSolvePolyAlgorithm, and does not touch unrelated differentiable types.Verified locally (Julia 1.11, NonlinearSolveBase dev'd from this branch +
Enzyme/ChainRulesCoreto trigger the extension):🤖 Generated with Claude Code