supports_initialization: use get_jumps to keep Enzyme readonly#4551
Conversation
`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>
Layer 2 status updateConfirmed independently that landing this Layer 1 fix is necessary but not sufficient: with this PR applied, bare
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 I also filed SciML/FunctionWrappersWrappers.jl#51 (defensive |
|
Test? I feel like this is very easy to accidentally break |
|
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. |
Please ignore until reviewed by @ChrisRackauckas.
Summary
supports_initialization(sys)short-circuits onisempty(get_systems(sys)), so the call tojumps(sys)here only ever reaches the early-return branch wherejumps(sys) === get_jumps(sys). However the unreachedcopy(js) + append!(js, namespace_jumps(subsys))branch injumpsis enough to prevent Enzyme's static analysis from provingsysread-only at this call site — and that breaksEnzyme.gradientthrough 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:
Before this PR:
EnzymeMutabilityExceptionatsupports_initialization→jumps(sys).After this PR:
Enzyme.gradientcompletes (returns a vector). The returned gradient itself is currently[0.0, …]because MTK#4181 (Enzyme rule forSetInitialUnknowns) 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.jlMixedDuplicated{NonlinearSolution}MethodError) still bite once you go pastremaketoremake + solve. Those need separate upstream work. The@test_broken Enzyme.gradientblocks in SciMLSensitivity'smtk.jl/parameter_initialization.jl/desauty_dae_mwe.jlstay until all four are resolved.Refs
SetInitialUnknowns— Layer 4, separate)Test plan
Enzyme.gradienton the MWE no longer errors