Enzyme reverse mode: wrong gradient for IIP wrapped function when args are mutated after the call

[ChrisRackauckas-Claude]

Summary

Reverse-mode Enzyme through a FunctionWrappersWrapper around an in-place function returns a silently wrong gradient when the caller mutates the wrapped function's arguments after the call — which is exactly what an ODE integrator does (f!(du,u,p,t) then step u). A single isolated call is correct; the bug only appears once an argument is mutated afterwards.

This is the root cause of the EnsembleProblem Enzyme adjoint failure in SciML/SciMLSensitivity.jl#1424 (under the default AutoSpecialize, DiffEqBase.wrapfun_iip wraps the ODE rhs in a FunctionWrappersWrapper, and the whole-solve Enzyme adjoint steps u between rhs calls). In the full nested ensemble it manifests as the Julia-1.11 _dispatch_ensemble_solve GC-root-rewrite segfault; reduced, it is a clean wrong-gradient.

MWE (self-contained)

using FunctionWrappersWrappers: FunctionWrappersWrapper
using Enzyme, FiniteDiff
using Enzyme: Reverse, set_runtime_activity

f!(du,u,p,t) = (du[1] = p[1]*u[1]; du[2] = p[2]*u[2]^2; nothing)
const ARGT = Tuple{Vector{Float64},Vector{Float64},Vector{Float64},Float64}

function loss(p)
    u = [1.5, 2.0]; du = zero(u)
    wf = FunctionWrappersWrapper(f!, (ARGT,), (Nothing,))
    wf(du, u, p, 0.0)         # du = f!(u,p) through the wrapper
    @. u = u + 0.05*du        # integrator-style mutation of u AFTER the call
    du[1]^2 + du[2]^2         # loss reads du only
end

p = [0.7, 0.4]
Enzyme.gradient(set_runtime_activity(Reverse), loss, p)  # -> [3.26, 13.84]   WRONG
FiniteDiff.finite_difference_gradient(loss, p)           # -> [3.15, 12.80]   correct
# deleting the `@. u = u + 0.05*du` line makes Enzyme correct

Verified on Julia 1.11.9, Enzyme 0.13.172/0.13.173, FunctionWrappersWrappers 1.9.3.

The wrong value is explainable exactly: the rule evaluates ∂du/∂p at the mutated u[1] = 1.5 + 0.05*du[1] = 1.5525 instead of the call-time 1.5, so ∂loss/∂p1 = 2*du1*u1 = 2*1.05*1.5525 = 3.26 instead of 2*1.05*1.5 = 3.15.

Root cause

In ext/FunctionWrappersWrappersEnzymeExt.jl, the Const-return (IIP) reverse path:

• EnzymeRules.augmented_primal(..., RT::Type{<:Const}, args...) runs the primal and returns AugmentedReturn(nothing, nothing, nothing) — no tape.
• EnzymeRules.reverse(..., dret::Type{<:Const}, tape, args...) recomputes the VJP with Enzyme.autodiff(Reverse, Const(f_orig), Const, args...) — i.e. against the arguments' current .vals.

Because nothing snapshots the call-time argument state, any mutation of an argument between the forward call and the reverse pass makes the recomputation differentiate f_orig about the wrong state. An integrator mutates u on every step, so every per-step gradient is wrong.

Fix

Snapshot the call-time argument values in augmented_primal's tape, and in reverse temporarily restore them before recomputing the VJP (then put the live values back). PR incoming.