EnzymeMutabilityException differentiating NonlinearSolve.solve over a RuntimeGeneratedFunction problem (Julia 1.12)

[ChrisRackauckas-Claude]

Summary

Enzyme.gradient over NonlinearSolve.solve works for a NonlinearProblem whose f is a plain closure / anonymous function, but fails with EnzymeMutabilityException as soon as f is a RuntimeGeneratedFunction (RGF) — same call shape, only the function wrapper differs.

This is the second layer of the Enzyme + ModelingToolkit remake stack documented in SciML/ModelingToolkit.jl#4550 / #4551 and SciML/SciMLSensitivity.jl#1323 / #1359. Layer 1 (MTK supports_initialization) is fixed by ModelingToolkit#4551. Layer 2 is reachable independently with this RGF-only MWE — no MTK in the loop, no SciMLSensitivity.

MWE (no MTK)

using RuntimeGeneratedFunctions, NonlinearSolve, Enzyme
using SciMLBase: NonlinearProblem, remake

RuntimeGeneratedFunctions.init(@__MODULE__)
f_rgf = @RuntimeGeneratedFunction(:((du, u, p) -> begin
    du[1] = u[1] - p[1] + p[2]; nothing
end))

# Plain closure: works
g_plain(du, u, p) = (du[1] = u[1] - p[1] + p[2]; nothing)
prob_plain = NonlinearProblem(g_plain, [0.0], [2.0, 1.0])
loss_plain(p) = sum(solve(remake(prob_plain; p)).u)
Enzyme.gradient(Enzyme.Reverse, loss_plain, [2.0, 1.0])  # PASS — [1.0, -1.0]

# Same shape with RGF: fails
prob = NonlinearProblem(f_rgf, [0.0], [2.0, 1.0])
loss_rgf(p) = sum(solve(remake(prob; p)).u)
Enzyme.gradient(Enzyme.Reverse, loss_rgf, [2.0, 1.0])

EnzymeMutabilityException: Function argument passed to autodiff cannot be proven readonly.
The potentially writing call is store i8 %"...iprob.f1.i.sroa.0.0.extract.trunc", ptr addrspace(11) %14, ...

Verified on Julia 1.12.6, Enzyme 0.13.148, RuntimeGeneratedFunctions 0.5.19, NonlinearSolve 4.19.1, SciMLBase 3.14.0.

What I ruled out

• Not an RGF body issue alone. Plain Enzyme.gradient(Reverse, p -> rgf(p[1]), [3.0]) works fine — RuntimeGeneratedFunction is differentiable in isolation. The failure mode is specifically RGF inside a NonlinearProblem under solve.

• Not opaque_closures = true. Reproducing with ModelingToolkit using eval_expression = true (so f is a plain var"#2#3" anonymous Julia function wrapped in GeneratedFunctionWrapper, no RGF) fails identically. The trigger is the wrapped-function pattern combined with solve, not RGF specifically.

• Not FunctionWrappersWrappers.SingleCacheStorage. Declaring EnzymeRules.inactive_type(::Type{SingleCacheStorage}) = true (filed as a defensive PR at Declare cache-storage types as Enzyme inactive SciML/FunctionWrappersWrappers.jl#51) does not change the outcome.

• set_runtime_activity(Reverse) does not help.

• Enzyme.Const(loss) advances past this point but then trips EnzymeRuntimeActivityError with "Mismatched activity for Unknown object of type System" inside SciMLBase's remake path — a separate failure layer.

Reduction further

The failing IR variable is iprob.f1.i.sroa.0.0, which SROAs through iprob.f (NonlinearFunction) into a sub-field of the wrapped function. The "store i8" is a Union-tag byte. The mutation is reported inside NonlinearProblem's mutable-struct constructor at SciMLBase/src/problems/nonlinear_problems.jl:185 — but it surfaces only when the function passed in is a wrapped function (RGF or otherwise), not when it's a plain closure.

I have not been able to reduce this to a pure Enzyme + plain Julia case yet — every wrapped-function reduction I tried still routes through NonlinearSolve. So the suspected interaction is between Enzyme's IPA on NonlinearProblem (mutable struct) and the way NonlinearSolve.solve dispatches on wrapped-function fields.

What this unblocks if fixed

Enzyme.gradient(Reverse, loss, tunables) through remake(prob; p = repack(tunables)) + solve(...) on any MTK-generated problem. Currently @test_broken in SciML/SciMLSensitivity.jl's test/mtk.jl, test/parameter_initialization.jl, test/desauty_dae_mwe.jl. Once this lands, ModelingToolkit#4551 (Layer 1) plus this fix flip those three blocks from @test_broken to passing.

Environment

• Julia 1.12.6
• Enzyme 0.13.148
• NonlinearSolve 4.19.1
• RuntimeGeneratedFunctions 0.5.19
• SciMLBase 3.14.0
• FunctionWrappersWrappers 1.9.0
• ModelingToolkit 11.26.5 (not required for the MWE, just for the user scenario)

[wsmoses]

EnzymeMutability Exception is a user error not an Enzyme error

@ChrisRackauckas

[ChrisRackauckas-Claude]

Acknowledged — EnzymeMutabilityException is correct behavior when a closure captures a mutable struct and the user can't annotate the capture. Sorry for the noise.

Re-ran the chain with the proper user-side annotations and reduced to a genuinely upstream-looking error one level down:

using RuntimeGeneratedFunctions, NonlinearSolve, Enzyme
using SciMLBase: NonlinearProblem, remake

RuntimeGeneratedFunctions.init(@__MODULE__)
f_rgf = @RuntimeGeneratedFunction(:((du, u, p) -> (du[1] = u[1] - p[1] + p[2]; nothing)))
prob = NonlinearProblem(f_rgf, [0.0], [2.0, 1.0])

# Explicit NewtonRaphson (avoid polyalgorithm Union → IllegalTypeAnalysis)
loss(p) = sum(solve(remake(prob; p), NewtonRaphson()).u)

# Const(loss) (capture is non-derivative) + set_runtime_activity for the
# downstream activity check
Enzyme.gradient(
    Enzyme.set_runtime_activity(Enzyme.Reverse),
    Enzyme.Const(loss),
    [2.0, 1.0]
)

That now reaches EnzymeRuntimeException inside DifferentiationInterfaceForwardDiffExt's prepare_jacobian_nokwarg with:

Enzyme: jl_call calling convention not implemented in aug_forward for
%jl_f__compute_sparams_ret = call nonnull "enzyme_type"="{[-1]:Pointer}" ptr addrspace(10) (ptr, ptr addrspace(10), ...) @julia.call(ptr nonnull @jl_f__compute_sparams, ...)

I'll close this issue and open a focused one on that jl_call-in-aug_forward path if you want it tracked separately — let me know if that's the right venue or if the actual fix lives on the DifferentiationInterface / NonlinearSolve side.

[wsmoses]

cc @gdalle @adrhill as an fyi this looks like a type stability potential problem in DI's fwddiff backend [should still fix the latter case here, but fyi]

[gdalle]

ForwardDiff's preparation is type-unstable by design due to automatic chunk size selection (unless a fixed chunk size is specified in the backend object). Is that the issue here?

[adrhill]

It likely is: JuliaDiff/DifferentiationInterface.jl#1020 (comment)