Skip to content

NonlinearSolveBase: mark nonlinear-solve algorithms Enzyme-inactive - #947

Merged
ChrisRackauckas merged 1 commit into
SciML:masterfrom
ChrisRackauckas-Claude:claude/enzyme-inactive-type-nonlinearalg
May 31, 2026
Merged

NonlinearSolveBase: mark nonlinear-solve algorithms Enzyme-inactive#947
ChrisRackauckas merged 1 commit into
SciML:masterfrom
ChrisRackauckas-Claude:claude/enzyme-inactive-type-nonlinearalg

Conversation

@ChrisRackauckas-Claude

Copy link
Copy Markdown
Member

Please ignore until reviewed by @ChrisRackauckas. Draft, opened by an agent.

What

Declare the whole AbstractNonlinearSolveAlgorithm hierarchy Enzyme.EnzymeRules.inactive_type, in NonlinearSolveBaseEnzymeExt:

Enzyme.EnzymeRules.inactive_type(::Type{<:NonlinearSolveBase.AbstractNonlinearSolveAlgorithm}) = true

Why

When a NonlinearProblem is solved under Enzyme.set_runtime_activity(Reverse) — most notably the ModelingToolkit DAE-initialization solve reached through get_initial_valuessolve_up — the default NonlinearSolvePolyAlgorithm arrives at the solve_up custom Enzyme rule as a positional args.... Nonlinear-solve algorithms are pure solver configuration (Vals, Rationals, Nothing/Missing sentinels, nested immutable sub-algorithms, and a ForwardDiff tag type parameter) with no differentiable floating-point data, but without an inactivity declaration set_runtime_activity promotes that configuration argument to Duplicated. That manifests, nondeterministically, as either:

  • the roots_activep (DFT_CONSTANT) != activep (DFT_DUP_ARG) assertion inside Enzyme's enzyme_custom_setup_args, or
  • a spurious algorithm shadow that corrupts activity bookkeeping for the genuinely-active u0/p, silently dropping a cotangent component.

Marking the algorithm hierarchy inactive forces Const, so activep == roots_activep == DFT_CONSTANT. This mirrors the established, shipping pattern in SciMLBase (inactive_type(::Type{<:AbstractSensitivityAlgorithm})) and LinearSolve (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-loaded EnzymeRules.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_autodiff wrap-skip and the inactive_kwarg for the solve_up rule are already merged). A downstream SciMLSensitivity PR will flip the corresponding mtk.jl test on and bump its lower bound once this is tagged.

Test

Adds lib/NonlinearSolveBase/test/enzyme_inactive_algorithm.jl (wired into runtests.jl), a deterministic regression test that the declaration covers AbstractNonlinearSolveAlgorithm, the concrete NonlinearSolvePolyAlgorithm, and does not touch unrelated differentiable types.

Verified locally (Julia 1.11, NonlinearSolveBase dev'd from this branch + Enzyme/ChainRulesCore to trigger the extension):

NonlinearSolveBase 2.30.2
Test Summary:                                                       | Pass  Total  Time
Enzyme inactive_type covers the nonlinear-solve algorithm hierarchy |    4      4  0.2s

🤖 Generated with Claude Code

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
ChrisRackauckas marked this pull request as ready for review May 31, 2026 18:31
@ChrisRackauckas
ChrisRackauckas merged commit 2d2ba94 into SciML:master May 31, 2026
91 of 121 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants