Snapshot call-time args in the IIP Enzyme reverse rule (fixes #62)

[ChrisRackauckas-Claude]

Please ignore until reviewed by @ChrisRackauckas. Draft, opened by an automated agent.

Fixes #62. Root cause and reduction of the EnsembleProblem Enzyme adjoint blocker in SciML/SciMLSensitivity.jl#1424.

Problem

The Const-return (IIP) reverse rule re-runs Enzyme.autodiff(Reverse, Const(f_orig), Const, args...) during the reverse pass, reading the arguments' current .vals, while augmented_primal saved no tape. When a caller mutates an argument after the forward call — an ODE integrator stepping u after every rhs call — the VJP is taken about the post-mutation state, so the gradient is silently wrong. A single isolated call is fine; the bug needs a mutation-after-call. In the full nested ensemble adjoint this surfaces on Julia 1.11 as the _dispatch_ensemble_solve GC-root-rewrite segfault (#1424).

Fix

• augmented_primal (Const return) snapshots the call-time argument values into the tape (before running the primal, in case f_orig mutates its inputs).
• reverse (Const return) temporarily restores those snapshots before the Enzyme.autodiff(Reverse, …) recompute, then puts the live values back so the surrounding reverse pass is undisturbed.
• _snapshot/_restore! helpers: copy/copyto! for arrays, no-op for immutables (scalars like t).

Only the IIP Const-return path changed; the forward, Active, and Duplicated/BatchDuplicated paths are untouched.

Verification (Julia 1.11.9, Enzyme 0.13.172/0.13.173, FunctionWrappersWrappers dev)

• The MWE from Enzyme reverse mode: wrong gradient for IIP wrapped function when args are mutated after the call #62 now matches FiniteDiff ([3.15, 12.80]); an 8-step integrator loop matches central differences.
• The existing test/Enzyme suite passes (one assertion updated: augmented_primal now tapes the arg snapshots instead of nothing).
• Two regression tests added: single wrapped IIP call with the argument mutated afterwards (analytic check), and a multi-step integrator (finite-difference check).
• End-to-end on the real EnsembleProblem adjoint (SciMLBase 3.30 / OrdinaryDiffEqTsit5 / Enzyme 0.13.172, set_runtime_activity(Reverse)): with this fix the default-AutoSpecialize solve goes from 4/4 deterministic SIGSEGV → no segfault, leaving only the separate with_logstate/AggregateLogger blocker (same residual as FullSpecialize).

🤖 Generated with Claude Code

[ChrisRackauckas-Claude]

Two notes after review feedback:

Scope — this is a safety net, not the primary path. SciMLSensitivity already unwrapped_fs the rhs before building every adjoint VJP (adjoint_common.jl, derivative_wrappers.jl, and the explicit sol.prob.f.f isa FunctionWrappersWrapper → unwrapped_f(...) rebuilds in interpolating_adjoint.jl/quadrature_adjoint.jl/gauss_adjoint.jl/backsolve_adjoint.jl), which is why the per-trajectory EnzymeVJP adjoint never hits this rule. The rule only matters when something differentiates the wrapper directly without unwrapping — e.g. an outer Enzyme.gradient straight over a whole solve(...)/solve(EnsembleProblem...). The complete fix for #1424 should also unwrap on the EnsembleProblem path (so outer Enzyme never sees the wrapper); this PR makes the wrapper correct when it is hit.

Allocation. The snapshot is gated on EnzymeRules.overwritten(config) (2nd commit): a non-mutating call copies nothing, and an integrator-style call copies only the overwritten arrays (du/u), not p or t. Since overwritten is a conservative over-approximation, gating can't drop a snapshot that's actually needed.

Verified: existing test/Enzyme suite green; new regression tests (single mutated-after-call + 8-step integrator) pass; and end-to-end the default-AutoSpecialize EnsembleProblem Enzyme adjoint goes from 4/4 SIGSEGV → no segfault with this change.