Skip to content

supports_initialization: use get_jumps to keep Enzyme readonly#4551

Merged
ChrisRackauckas merged 2 commits into
SciML:masterfrom
ChrisRackauckas-Claude:cc/supports-init-get-jumps
May 21, 2026
Merged

supports_initialization: use get_jumps to keep Enzyme readonly#4551
ChrisRackauckas merged 2 commits into
SciML:masterfrom
ChrisRackauckas-Claude:cc/supports-init-get-jumps

Conversation

@ChrisRackauckas-Claude
Copy link
Copy Markdown

Please ignore until reviewed by @ChrisRackauckas.

Summary

supports_initialization(sys) short-circuits on isempty(get_systems(sys)), so the call to jumps(sys) here only ever reaches the early-return branch where jumps(sys) === get_jumps(sys). However the unreached copy(js) + append!(js, namespace_jumps(subsys)) branch in jumps is enough to prevent Enzyme's static analysis from proving sys read-only at this call site — and that breaks Enzyme.gradient through any closure that reaches _late_binding_update_u0_p_impl, e.g. remake(prob; p = repack(tunables)) on any MTK-generated problem.

Swap to get_jumps(sys), which is a pure field accessor and identical in this short-circuited position.

Verification

The MWE from #4550:

@parameters a b c
@variables x(t) y(t) z(t)
eqs = [D(x) ~ a*x + y + z, 0 ~ y^3 + y - b*x, 0 ~ z^3 + z - c*y]
@mtkbuild sys = ODESystem(eqs, t)
prob = ODEProblem(sys, [x => 1.0], (0.0, 0.1),
                  [a => -0.5, b => 2.0, c => 1.5];
                  guesses = [y => 1.0, z => 0.5], use_scc = false)
iprob = prob.f.initialization_data.initializeprob
itunables, irepack, _ = SS.canonicalize(SS.Tunable(), parameter_values(iprob))
remake_only = let iprob = iprob, irepack = irepack
    p -> sum(remake(iprob, p = irepack(p)).u0)
end
Enzyme.gradient(Enzyme.Reverse, remake_only, itunables)

Before this PR: EnzymeMutabilityException at supports_initializationjumps(sys).

After this PR: Enzyme.gradient completes (returns a vector). The returned gradient itself is currently [0.0, …] because MTK#4181 (Enzyme rule for SetInitialUnknowns) is still missing — that's a separate semantic-correctness issue that this PR doesn't touch; it just removes the compile-time blocker.

What downstream this unblocks

This is Layer 1 of a 4-layer Enzyme-through-MTK-remake stack documented during the SciMLSensitivity #1323 / #1359 investigation. Layers 2–4 (FunctionWrappersWrappers cache mutation, SciMLBase Const(loss) EnzymeRuntimeActivityError, Enzyme.jl MixedDuplicated{NonlinearSolution} MethodError) still bite once you go past remake to remake + solve. Those need separate upstream work. The @test_broken Enzyme.gradient blocks in SciMLSensitivity's mtk.jl / parameter_initialization.jl / desauty_dae_mwe.jl stay until all four are resolved.

Refs

Test plan

  • Local: Enzyme.gradient on the MWE no longer errors
  • CI green

`supports_initialization(sys)` short-circuits on `isempty(get_systems(sys))`,
so the call to `jumps(sys)` only ever reaches the early-return branch
(`jumps(sys) === get_jumps(sys)` when there are no subsystems). However the
unreached `copy(js) + append!(js, namespace_jumps(subsys))` branch in
`jumps` is enough to prevent Enzyme's static analysis from proving `sys`
read-only, so any `Enzyme.gradient` over a closure that reaches
`_late_binding_update_u0_p_impl` (e.g. `remake(prob; p = repack(tunables))`
on MTK-generated problems) throws `EnzymeMutabilityException`.

`get_jumps` is a plain field accessor and equivalent here, so this swap
keeps semantics identical while unblocking `Enzyme.gradient` through
`remake` on MTK problems. Verified with the issue MWE — `Enzyme.gradient`
now completes instead of erroring. (The remaining MTK#4181 piece — adding
an Enzyme rule for `SetInitialUnknowns` so the returned gradient is
semantically correct — is independent of this fix.)

Closes SciML#4550.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
@ChrisRackauckas ChrisRackauckas marked this pull request as ready for review May 21, 2026 13:57
@ChrisRackauckas-Claude
Copy link
Copy Markdown
Author

Layer 2 status update

Confirmed independently that landing this Layer 1 fix is necessary but not sufficient: with this PR applied, bare Enzyme.gradient through remake succeeds, but remake + solve still fails with EnzymeMutabilityException. I reduced Layer 2 to a no-MTK MWE and filed it upstream at EnzymeAD/Enzyme.jl#3117:

Enzyme.gradient over NonlinearSolve.solve works for a NonlinearProblem whose f is a plain anonymous function, but fails as soon as f is a RuntimeGeneratedFunction (or any wrapped-function pattern — eval_expression = true still fails). The failing IR variable is iprob.f1.i.sroa.0.0 — a Union-tag byte store SROA'd through NonlinearFunction.f reported inside SciMLBase's mutable NonlinearProblem constructor.

The full breakdown of layers 2–4 is in the discussion on SciMLSensitivity#1359. This PR remains the right Layer 1 fix on its own; the three @test_broken Enzyme.gradient blocks in SciMLSensitivity (mtk.jl, parameter_initialization.jl, desauty_dae_mwe.jl) flip to passing once both this and EnzymeAD/Enzyme.jl#3117 land.

I also filed SciML/FunctionWrappersWrappers.jl#51 (defensive inactive_type rule for cache storage) — that one doesn't fix Layer 2 by itself but removes one conservative-IPA layer.

@AayushSabharwal
Copy link
Copy Markdown
Member

Test? I feel like this is very easy to accidentally break

@ChrisRackauckas
Copy link
Copy Markdown
Member

The only real answer for whether Enzyme works is if Enzyme works, so we'll get that in the integration tests at the end of the journey.

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.

Enzyme.gradient through remake fails with EnzymeMutabilityException via supports_initialization → jumps(sys)

3 participants