Enzyme support for NonlinearProblem and SCCNonlinearProblem differentiation (#1358) by ChrisRackauckas-Claude · Pull Request #1397 · SciML/SciMLSensitivity.jl

ChrisRackauckas-Claude · 2026-03-27T19:17:19Z

Summary

Enable Enzyme differentiation through NonlinearProblem and SCCNonlinearProblem initialization solves, addressing #1358.

Changes

Full-parameter VJP in SteadyStateAdjoint (steadystate_adjoint.jl):
- The existing VJP only computed gradients w.r.t. tunable parameters, missing caches/non-tunable components
- For SCCNonlinearProblem where explicitfuns! write active data into caches, this caused zero gradients
- Added a full-parameter Zygote pullback that captures ALL parameter dependencies (tunables + caches)
- Returns NamedTuple gradient for EnzymeOriginator so Enzyme can chain through explicitfuns! mutations
EnzymeOriginator tangent handling (concrete_solve.jl):
- For EnzymeOriginator with SciMLStructure params, pass the full parameter gradient through (not just tunable vector)
- Enables the Enzyme reverse rule to accumulate into all shadow components
Documentation (test files):
- Document specific Enzyme blockers in @test_broken annotations

Companion PRs

NonlinearSolve.jl#884: SCC simplifications + NamedTuple field accumulation in Enzyme reverse rule

Verified

NonlinearProblem + explicit alg + MTKParameters: gradient matches FiniteDiff ✓
explicitfun! → caches → sub-problem solve (SCC pattern): gradient matches FiniteDiff ✓

Test plan

CI passes
Existing NonlinearProblem adjoint tests still pass (only EnzymeOriginator path changed)

🤖 Generated with Claude Code

ChrisRackauckas-Claude · 2026-03-28T20:38:30Z

Changes in latest commit

1. Unified `automatic_sensealg_choice` signatures

Added original_p parameter to the ODE dispatch (AbstractODEProblem/AbstractSDEProblem), matching the ConcreteNonlinearProblem dispatch signature
Both dispatches now detect non-tunable active components (caches) and set diff_tunables=Val(false) accordingly

2. Added `diff_tunables` to `GaussAdjoint` and `QuadratureAdjoint`

Added diff_tunables::DT field (default Val(true)) to both structs, mirroring SteadyStateAdjoint
Updated constructors, setvjp, and all type parameter lists
Propagated use_full_p through:
- ODEGaussAdjointSensitivityFunction → adjointdiffcache
- ODEQuadratureAdjointSensitivityFunction → adjointdiffcache
- GaussIntegrand (tunables/repack setup)
- AdjointSensitivityIntegrand (tunables/repack setup)
- update_p_integrand (tunables/repack setup)
- _adjoint_sensitivities_gauss (tunables/repack setup)
Added diff_tunables handling in ODE backpass for Enzyme full-parameter gradients
All diff_tunables accesses guarded with hasproperty for GaussKronrodAdjoint compatibility

3. Removed `scc_enzyme.jl`

scc_nonlinearsolve.jl is already in runtests.jl Core 8 group

Testing

Verified constructors work with diff_tunables=Val(false) and Val(true) (default)
setvjp preserves diff_tunables
No FieldError regressions (previously crashed on GaussKronrodAdjoint)
Core8 crash is pre-existing LLVM/Enzyme issue on Julia 1.12 (not related to these changes)

…L#1358) Add detailed comments to test files and source code explaining the specific upstream blockers preventing Enzyme from differentiating through NonlinearProblem/SCCNonlinearProblem initialization: - Julia 1.12: LLVM crash due to NonlinearSolve Enzyme rules disabled - Julia 1.10: MixedReturnException, EnzymeMutabilityException, NamedTuple broadcasting errors in reverse rule - SCCNonlinearProblem: no _concrete_solve_adjoint dispatch exists References: NonlinearSolve.jl#869, Enzyme.jl#2699 Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

…#1358) For EnzymeOriginator with SciMLStructure parameters, return the tunable gradient vector directly from steadystatebackpass instead of the Zygote-repacked NamedTuple. The NonlinearSolveBaseEnzymeExt reverse rule uses SciMLStructures.replace! to accumulate it into the parameter shadow, going through the proper SciMLStructures interface. This avoids the NamedTuple broadcasting error and ensures all tangent accumulation uses SciMLStructures.canonicalize/replace! rather than making assumptions about the NamedTuple field structure. Companion PR: NonlinearSolve.jl#879 Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Add `diff_tunables::Val` field to SteadyStateAdjoint (default Val(true)). When Val(false), the parameter VJP via vecjacobian! is computed w.r.t. the full parameter object (including caches) instead of just tunables. This is needed for SCCNonlinearProblem where explicitfuns! write active data into non-tunable parameter components (caches). The automatic sensealg choice detects non-empty caches and sets diff_tunables=Val(false) with a structured-VJP-compatible backend. Changes: - sensitivity_algorithms.jl: Add diff_tunables field to SteadyStateAdjoint - adjoint_common.jl: Add use_full_p kwarg to adjointdiffcache - steadystate_adjoint.jl: Gate use_full_p on diff_tunables flag - concrete_solve.jl: Pass original_p to automatic_sensealg_choice for cache detection; return full gradient for EnzymeOriginator when diff_tunables=Val(false) - test/scc_enzyme.jl: Direct SCC differentiation test with Enzyme Companion PR: NonlinearSolve.jl#884 Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Test SCC with direct construction (2-component problem with parameter coupling through caches): - FiniteDiff: passes (ground truth) - Mooncake: passes - ForwardDiff: @test_broken (Dual numbers into Float64 mutation buffer) - Enzyme: @test_broken (EnzymeNoTypeError in SCC dispatch chain) Added to runtests.jl Core 8 group. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Mooncake's build_rrule on the MTK DAE initialization problem takes too long for CI timeouts. SCCNonlinearProblem Mooncake test is in scc_nonlinearsolve.jl instead. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

…aussAdjoint/QuadratureAdjoint - Add original_p parameter to ODE automatic_sensealg_choice dispatch, matching the NonlinearProblem dispatch signature - Add diff_tunables::Val field to GaussAdjoint and QuadratureAdjoint (default Val(true)), mirroring SteadyStateAdjoint - Propagate use_full_p through GaussIntegrand, AdjointSensitivityIntegrand, and adjointdiffcache for both adjoint methods - Handle diff_tunables in ODE backpass for Enzyme full-parameter gradients - Guard all diff_tunables access with hasproperty for GaussKronrodAdjoint compatibility (which shares GaussAdjoint code paths) - Remove scc_enzyme.jl (scc_nonlinearsolve.jl already in runtests.jl) Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

- Add SCCNonlinearSolve to [extras] and [targets] in Project.toml (was missing, causing LoadError on CI for scc_nonlinearsolve.jl) - Fix runic formatting: add explicit return statements in scc_nonlinearsolve.jl functions Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

With Mooncake v0.5.25 + ModelingToolkitBase v1.28.0, the non-SCC init path works. The SCC path remains @test_broken. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

…ion) Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

…pport) The full-parameter VJP path (when diff_tunables=Val(false)) was gated on EnzymeOriginator only, defeating the purpose of diff_tunables for ChainRules-based AD backends like Mooncake. This blocked SCC gradient flow for Mooncake even after SciML#1397. Fix: extend the full-p path to all originators. Convert dp_full from its native SciMLStructure form (e.g. MTKParameters) to a NamedTuple for non-Enzyme originators, since Mooncake's @from_chainrules bridge expects NamedTuple tangents while Enzyme accumulates into the struct shadow directly. Verified: Mooncake through SCCNonlinearProblem init in the De Sauty DAE test now produces gradients matching FiniteDiff for use_scc=true. non-SCC path still works. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

ChrisRackauckas-Claude force-pushed the enzyme-1358-test-improvements branch 4 times, most recently from b007f44 to 186d502 Compare March 28, 2026 11:04

ChrisRackauckas-Claude closed this Mar 28, 2026

ChrisRackauckas-Claude reopened this Mar 28, 2026

ChrisRackauckas-Claude mentioned this pull request Apr 8, 2026

Add diff_tunables to SteadyStateAdjoint for full-parameter VJP #1407

Closed

2 tasks

ChrisRackauckas and others added 8 commits April 8, 2026 09:03

Add SCCNonlinearSolve compat entry to fix QA/Aqua test

ad3951c

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

ChrisRackauckas-Claude force-pushed the enzyme-1358-test-improvements branch from 0595d4e to ad3951c Compare April 8, 2026 13:03

ChrisRackauckas and others added 6 commits April 8, 2026 09:05

Bump FunctionWrappersWrappers compat to include v1.2 (Mooncake extens…

09eb736

…ion) Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Retrigger CI after FunctionWrappersWrappers v1.2.0 registration

2213ffd

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Retrigger CI after FunctionWrappersWrappers v1.2.1 fix

c155806

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Retrigger CI after FunctionWrappersWrappers v1.4.0 with Mooncake fix

a5f51c3

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Require FunctionWrappersWrappers >= 1.4 for working Mooncake extension

0ce8d64

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

ChrisRackauckas merged commit 09d2d53 into SciML:master Apr 10, 2026
46 of 50 checks passed

This was referenced Apr 10, 2026

Enable full-parameter VJP for Mooncake (SCCNonlinearProblem support) #1412

Merged

[WIP] docs: prefer Mooncake over Zygote where it works end-to-end #1419

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Enzyme support for NonlinearProblem and SCCNonlinearProblem differentiation (#1358)#1397

Enzyme support for NonlinearProblem and SCCNonlinearProblem differentiation (#1358)#1397
ChrisRackauckas merged 14 commits into
SciML:masterfrom
ChrisRackauckas-Claude:enzyme-1358-test-improvements

ChrisRackauckas-Claude commented Mar 27, 2026

Uh oh!

ChrisRackauckas-Claude commented Mar 28, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

ChrisRackauckas-Claude commented Mar 27, 2026

Summary

Changes

Companion PRs

Verified

Test plan

Uh oh!

ChrisRackauckas-Claude commented Mar 28, 2026

Changes in latest commit

1. Unified automatic_sensealg_choice signatures

2. Added diff_tunables to GaussAdjoint and QuadratureAdjoint

3. Removed scc_enzyme.jl

Testing

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

1. Unified `automatic_sensealg_choice` signatures

2. Added `diff_tunables` to `GaussAdjoint` and `QuadratureAdjoint`

3. Removed `scc_enzyme.jl`