From a780e767de57abac3b5fd2914731a4a069227684 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Mon, 29 Jun 2026 12:38:49 -0400 Subject: [PATCH 1/3] Fix Enzyme reverse rule: snapshot call-time args for IIP wrapped functions The Const-return (IIP) reverse rule re-runs `Enzyme.autodiff(Reverse, ...)` on the unwrapped function during the reverse pass, reading the arguments' current `.val`s. When a caller mutates those arguments after the forward call -- e.g. an ODE integrator stepping `u` after every RHS call -- the VJP was taken about the post-mutation state, giving a silently wrong gradient. In the full nested EnsembleProblem adjoint this surfaced as the Julia-1.11 `_dispatch_ensemble_solve` GC-root-rewrite segfault (SciMLSensitivity.jl#1424). `augmented_primal` now snapshots the call-time argument values into the tape, and `reverse` temporarily restores them before recomputing the VJP (then puts the live values back so the surrounding reverse pass is undisturbed). Adds regression tests: a single wrapped IIP call with the argument mutated afterwards, and an 8-step integrator loop, both checked against analytic / finite-difference references. Co-Authored-By: Chris Rackauckas Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01BZwV7CBdLPqVpJrvRAWcv6 --- ext/FunctionWrappersWrappersEnzymeExt.jl | 32 +++++++++- test/Enzyme/enzyme_tests.jl | 78 ++++++++++++++++++++++-- 2 files changed, 103 insertions(+), 7 deletions(-) diff --git a/ext/FunctionWrappersWrappersEnzymeExt.jl b/ext/FunctionWrappersWrappersEnzymeExt.jl index 1c05580..c5157ec 100644 --- a/ext/FunctionWrappersWrappersEnzymeExt.jl +++ b/ext/FunctionWrappersWrappersEnzymeExt.jl @@ -154,6 +154,18 @@ end # Reverse mode rules # ============================================================================= +# The Const-return (IIP) reverse rule re-runs `Enzyme.autodiff(Reverse, …)` on +# the unwrapped function during the *reverse* pass. That recomputation reads +# the arguments' current `.val`s — but a caller may mutate those arguments +# between the forward call and the reverse pass (an ODE integrator steps `u` +# after every RHS call). Without snapshotting, the VJP is then taken about the +# wrong state and the gradient is silently wrong. `_snapshot`/`_restore!` let +# the rule pin the call-time argument values in its tape. +@inline _snapshot(v::AbstractArray) = copy(v) +@inline _snapshot(v) = v +@inline _restore!(dst::AbstractArray, snap::AbstractArray) = (copyto!(dst, snap); nothing) +@inline _restore!(@nospecialize(dst), @nospecialize(snap)) = nothing + function EnzymeRules.augmented_primal( config::EnzymeRules.RevConfig, func::EnzymeCore.Annotation{<:FunctionWrappersWrapper}, @@ -172,8 +184,12 @@ function EnzymeRules.augmented_primal( end # Const return (e.g. IIP functions returning Nothing, or any non-differentiated -# return). Just run the primal for its side effects; no tape is needed because -# the reverse pass has nothing to propagate back from the return. +# return). Run the primal for its side effects, and snapshot the call-time +# argument values into the tape so the reverse pass can take the VJP about the +# state at the time of *this* call — even if the caller mutates the arguments +# afterwards (the ODE-integrator pattern: `wf(du, u, p, t)` followed by an +# in-place step on `u`). Snapshot BEFORE running the primal in case `f_orig` +# mutates its inputs. function EnzymeRules.augmented_primal( config::EnzymeRules.RevConfig, func::EnzymeCore.Annotation{<:FunctionWrappersWrapper}, @@ -181,9 +197,10 @@ function EnzymeRules.augmented_primal( args::Vararg{EnzymeCore.Annotation, N} ) where {N} f_orig = unwrap(func.val) + tape = ntuple(i -> _snapshot(args[i].val), Val(N)) pargs = ntuple(i -> args[i].val, Val(N)) f_orig(pargs...) - return EnzymeRules.AugmentedReturn(nothing, nothing, nothing) + return EnzymeRules.AugmentedReturn(nothing, nothing, tape) end # Duplicated / BatchDuplicated return: record the primal so that reverse has @@ -315,7 +332,16 @@ function EnzymeRules.reverse( # Only worth invoking Enzyme.autodiff when at least one arg is # Duplicated/BatchDuplicated — otherwise there's nothing to accumulate. if any(a -> a isa EnzymeCore.Duplicated || a isa EnzymeCore.BatchDuplicated, args) + # Temporarily restore the call-time argument values (snapshotted in the + # tape during augmented_primal) so the VJP is taken about the correct + # state, then put the live values back so the surrounding reverse pass + # is left undisturbed. Without this, a caller that mutates an argument + # after the forward call (an ODE integrator stepping `u`) makes the + # recomputation differentiate `f_orig` about the wrong state. + live = ntuple(i -> _snapshot(args[i].val), Val(N)) + map((a, snap) -> _restore!(a.val, snap), args, tape) Enzyme.autodiff(Reverse, Const(f_orig), Const, args...) + map((a, snap) -> _restore!(a.val, snap), args, live) end return ntuple(Val(N)) do i if args[i] isa EnzymeCore.Active diff --git a/test/Enzyme/enzyme_tests.jl b/test/Enzyme/enzyme_tests.jl index 5a8f5f9..37973b4 100644 --- a/test/Enzyme/enzyme_tests.jl +++ b/test/Enzyme/enzyme_tests.jl @@ -164,9 +164,11 @@ end @testset "Enzyme reverse mode, Const return — augmented_primal runs primal" begin # Mirrors the forward {false, false} case on the reverse side. Augmented - # primal runs the wrapped function for its side effects and returns - # AugmentedReturn(nothing, nothing, nothing). Reverse returns `nothing` - # per arg since there is no return derivative to propagate. + # primal runs the wrapped function for its side effects and tapes a snapshot + # of the call-time argument values (so a later mutation by the caller can't + # make the reverse pass differentiate about the wrong state). Reverse + # returns `nothing`/zero per arg since there is no return derivative to + # propagate. counter = Ref(0) g(x, y) = (counter[] += 1; x + y) # returns Float64 (ignored via Const RT) fww = FunctionWrappersWrapper(g, (Tuple{Float64, Float64},), (Float64,)) @@ -183,7 +185,7 @@ end @test counter[] == 1 # primal ran exactly once @test aug.primal === nothing # NeedsPrimal == false @test aug.shadow === nothing - @test aug.tape === nothing + @test aug.tape == (3.0, 4.0) # call-time snapshots of the args # Reverse step — dret is Const (passed as TYPE not instance in reverse # rules). Enzyme's rule protocol requires concrete gradients for Active @@ -330,6 +332,74 @@ end @test u_shadow[1] ≈ expected_u_grad end +# ============================================================================= +# Regression for the wrong gradient when a wrapped IIP function's arguments are +# MUTATED AFTER the call — the ODE-integrator pattern that the whole-solve +# Enzyme adjoint exercises (and the root cause of the EnsembleProblem adjoint +# failure, SciMLSensitivity.jl#1424). +# +# The Const-return reverse rule re-runs `Enzyme.autodiff(Reverse, …)` on the +# unwrapped function during the reverse pass. If it differentiates about the +# arguments' *current* state rather than their *call-time* state, then any +# caller that steps `u` after the RHS call gets a silently wrong gradient. +# Before the snapshot/restore tape fix these end-to-end gradients were wrong. +# ============================================================================= + +@testset "Enzyme Reverse: IIP wrapper, args mutated after call (single step)" begin + f!(du, u, p, t) = (du[1] = p[1] * u[1]; du[2] = p[2] * u[2]^2; nothing) + 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) + @inbounds for k in 1:2 + u[k] += 0.05 * du[k] # mutate u AFTER the wrapped call + end + return du[1]^2 + du[2]^2 # loss depends on du only + end + + p = [0.7, 0.4] + # du = [p1*1.5, p2*4]; loss = (1.5 p1)^2 + (4 p2)^2 + # ∂loss/∂p = [2*1.5^2*p1, 2*4^2*p2] = [4.5 p1, 32 p2]; evaluated at CALL-TIME u + g = collect(Enzyme.gradient(Enzyme.set_runtime_activity(Enzyme.Reverse), loss, p)[1]) + @test g ≈ [4.5 * p[1], 32 * p[2]] +end + +@testset "Enzyme Reverse: IIP wrapper in a multi-step integrator" begin + f!(du, u, p, t) = ( + du[1] = -p[1] * u[1] + p[2] * u[2]; + du[2] = -p[3] * u[2] + p[4] * u[1]; nothing + ) + ARGT = Tuple{Vector{Float64}, Vector{Float64}, Vector{Float64}, Float64} + + function loss(p) + u = [1.0, 2.0] + du = zero(u) + wf = FunctionWrappersWrapper(f!, (ARGT,), (Nothing,)) + for _ in 1:8 + wf(du, u, p, 0.0) + @inbounds for k in 1:2 + u[k] += 0.05 * du[k] # integrator step mutates u each call + end + end + return sum(abs2, u) + end + + p = [1.0, 0.5, 2.0, 0.3] + g = collect(Enzyme.gradient(Enzyme.set_runtime_activity(Enzyme.Reverse), loss, p)[1]) + + # central finite-difference reference (no extra deps) + fd = map(eachindex(p)) do i + h = 1.0e-6 + pp = copy(p); pp[i] += h + pm = copy(p); pm[i] -= h + (loss(pp) - loss(pm)) / (2h) + end + @test g ≈ fd rtol = 1.0e-4 +end + # ============================================================================= # Runtime-activity propagation through the FWW forward rules. # From d8f26a1f2243b9287a1171a7dfcec41355a0d26f Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Mon, 29 Jun 2026 13:25:17 -0400 Subject: [PATCH 2/3] Gate the IIP reverse-rule snapshot on overwritten() to minimize allocation The snapshot/restore safety net only needs to copy arguments that are actually mutated between the forward and reverse passes. Gate it on `EnzymeRules.overwritten(config)` (a conservative over-approximation, so it can't miss a needed snapshot): a non-mutating call now copies nothing, and an integrator-style call copies only the overwritten arrays (`du`/`u`), not `p` or the scalar `t`. Also fixes two pre-existing hand-built `RevConfig`s in the tests whose `Overwritten` tuples were one entry short (it is indexed `(func, args...)`). Co-Authored-By: Chris Rackauckas Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01BZwV7CBdLPqVpJrvRAWcv6 --- ext/FunctionWrappersWrappersEnzymeExt.jl | 13 +++++++++++-- test/Enzyme/enzyme_tests.jl | 9 ++++++--- 2 files changed, 17 insertions(+), 5 deletions(-) diff --git a/ext/FunctionWrappersWrappersEnzymeExt.jl b/ext/FunctionWrappersWrappersEnzymeExt.jl index c5157ec..6e9b8e6 100644 --- a/ext/FunctionWrappersWrappersEnzymeExt.jl +++ b/ext/FunctionWrappersWrappersEnzymeExt.jl @@ -161,6 +161,13 @@ end # after every RHS call). Without snapshotting, the VJP is then taken about the # wrong state and the gradient is silently wrong. `_snapshot`/`_restore!` let # the rule pin the call-time argument values in its tape. +# +# To keep this cheap, only arguments that Enzyme's `overwritten(config)` +# analysis flags as possibly modified between the forward and reverse passes +# are copied; everything else (and immutables like `t`) tapes `nothing` and is +# never copied. `overwritten` is a conservative over-approximation, so gating +# on it can never *miss* an argument that actually needs snapshotting. A +# non-mutating call therefore allocates nothing here. @inline _snapshot(v::AbstractArray) = copy(v) @inline _snapshot(v) = v @inline _restore!(dst::AbstractArray, snap::AbstractArray) = (copyto!(dst, snap); nothing) @@ -197,7 +204,9 @@ function EnzymeRules.augmented_primal( args::Vararg{EnzymeCore.Annotation, N} ) where {N} f_orig = unwrap(func.val) - tape = ntuple(i -> _snapshot(args[i].val), Val(N)) + # `overwritten` is indexed (func, args...), so arg `i` is entry `i + 1`. + ow = EnzymeRules.overwritten(config) + tape = ntuple(i -> (ow[i + 1] ? _snapshot(args[i].val) : nothing), Val(N)) pargs = ntuple(i -> args[i].val, Val(N)) f_orig(pargs...) return EnzymeRules.AugmentedReturn(nothing, nothing, tape) @@ -338,7 +347,7 @@ function EnzymeRules.reverse( # is left undisturbed. Without this, a caller that mutates an argument # after the forward call (an ODE integrator stepping `u`) makes the # recomputation differentiate `f_orig` about the wrong state. - live = ntuple(i -> _snapshot(args[i].val), Val(N)) + live = ntuple(i -> (tape[i] === nothing ? nothing : _snapshot(args[i].val)), Val(N)) map((a, snap) -> _restore!(a.val, snap), args, tape) Enzyme.autodiff(Reverse, Const(f_orig), Const, args...) map((a, snap) -> _restore!(a.val, snap), args, live) diff --git a/test/Enzyme/enzyme_tests.jl b/test/Enzyme/enzyme_tests.jl index 37973b4..b08d25b 100644 --- a/test/Enzyme/enzyme_tests.jl +++ b/test/Enzyme/enzyme_tests.jl @@ -175,7 +175,9 @@ end # Construct a concrete RevConfig. Fields: # (NeedsPrimal, NeedsShadow, Width, Overwritten, RuntimeActivity, StrongZero) - rconfig = EnzymeRules.RevConfig{false, false, 1, (false, false), false, false}() + # Overwritten is indexed (func, args...) — here (func, x, y). Mark only `x` + # as overwritten so we can check the rule snapshots exactly that arg. + rconfig = EnzymeRules.RevConfig{false, false, 1, (false, true, false), false, false}() counter[] = 0 aug = EnzymeRules.augmented_primal( @@ -185,7 +187,7 @@ end @test counter[] == 1 # primal ran exactly once @test aug.primal === nothing # NeedsPrimal == false @test aug.shadow === nothing - @test aug.tape == (3.0, 4.0) # call-time snapshots of the args + @test aug.tape == (3.0, nothing) # only the overwritten arg (x) is snapshotted # Reverse step — dret is Const (passed as TYPE not instance in reverse # rules). Enzyme's rule protocol requires concrete gradients for Active @@ -504,7 +506,8 @@ end du = [0.0]; du_shadow = [1.0] u = [3.0]; u_shadow = [0.0] - rconfig = EnzymeRules.RevConfig{false, false, 1, (false, false), false, false}() + # Overwritten indexed (func, du, u); none modified between fwd and rev here. + rconfig = EnzymeRules.RevConfig{false, false, 1, (false, false, false), false, false}() aug = EnzymeRules.augmented_primal( rconfig, Duplicated(fww, fww), # <-- Duplicated FWW From df23e99441ce349a961307bbd68a2762f1a4e726 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Thu, 2 Jul 2026 13:07:07 -0400 Subject: [PATCH 3/3] Handle Active args in the IIP reverse rule (mixed Duplicated/Active) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Const-return reverse rule returned `ntuple(i -> args[i] isa Active ? zero(...) : nothing)`, which (1) was union-typed — Enzyme rejects a `Tuple{Union{Nothing,Float64},…}` return with `ReverseRuleReturnError`, it requires exact per-slot types — and (2) zeroed every Active arg's gradient. This broke any wrapped call with a mix of Duplicated and Active args, e.g. a time-dependent in-place rhs `f!(du, u, p, t)` where `t` flows in as Active. `Enzyme.autodiff(Reverse, Const(f_orig), Const, args...)[1]` already returns the per-argument gradient tuple in the exact form a reverse rule must hand back (`nothing` for Const/Duplicated, the real gradient for Active). Forward it, rebuilt through `map(_revslot, args, raw)` so dispatching on each argument's concrete annotation type gives an exactly-typed result (the raw tuple is `Any`-typed inside the rule). Adds a mixed Duplicated/Active regression test. Co-Authored-By: Chris Rackauckas Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01BZwV7CBdLPqVpJrvRAWcv6 --- ext/FunctionWrappersWrappersEnzymeExt.jl | 41 +++++++++++++++--------- test/Enzyme/enzyme_tests.jl | 27 ++++++++++++++++ 2 files changed, 52 insertions(+), 16 deletions(-) diff --git a/ext/FunctionWrappersWrappersEnzymeExt.jl b/ext/FunctionWrappersWrappersEnzymeExt.jl index 6e9b8e6..1a852ce 100644 --- a/ext/FunctionWrappersWrappersEnzymeExt.jl +++ b/ext/FunctionWrappersWrappersEnzymeExt.jl @@ -173,6 +173,16 @@ end @inline _restore!(dst::AbstractArray, snap::AbstractArray) = (copyto!(dst, snap); nothing) @inline _restore!(@nospecialize(dst), @nospecialize(snap)) = nothing +# Build one slot of a reverse rule's return tuple: `nothing` for every +# non-Active arg (their gradients accumulate in-place), and the concrete +# gradient for each `Active` arg. Dispatching on the annotation *type* keeps +# the resulting tuple exactly typed (e.g. `Tuple{Nothing, Nothing, Float64}`) +# even though the raw `autodiff`/thunk return is `Any`-typed inside the rule — +# Enzyme rejects a union-typed return. `g` is the matching entry of that raw +# per-argument gradient tuple. +@inline _revslot(::EnzymeCore.Active{T}, g) where {T} = convert(T, g)::T +@inline _revslot(::EnzymeCore.Annotation, @nospecialize(g)) = nothing + function EnzymeRules.augmented_primal( config::EnzymeRules.RevConfig, func::EnzymeCore.Annotation{<:FunctionWrappersWrapper}, @@ -325,11 +335,14 @@ end # accumulate into any `Duplicated` arg shadow buffers (the SciML IIP # pattern). Simply returning `nothing` left Duplicated shadows at zero. # -# Per Enzyme's rule return-type protocol, `Active` args require a concrete -# scalar gradient (not `nothing`). Under a `Const` return there is no -# gradient source, so Active arg gradients are zero. `Duplicated` / -# `BatchDuplicated` args return `nothing` because their gradients are -# accumulated in-place by the `Enzyme.autodiff(Reverse, …)` call above. +# `Enzyme.autodiff(Reverse, Const(f_orig), Const, args...)[1]` already returns +# the per-argument gradient tuple in exactly the form a reverse rule must hand +# back: `nothing` for each `Const`/`Duplicated`/`BatchDuplicated` arg (their +# gradients are accumulated in-place), and the concrete gradient for each +# `Active` arg — with exact per-slot types (Enzyme rejects a union-typed +# `Tuple{Union{Nothing,Float64},…}`). So we forward that tuple directly, which +# also makes `Active` args (e.g. `t` in a time-dependent IIP rhs) correct +# rather than zeroed. function EnzymeRules.reverse( config::EnzymeRules.RevConfig, func::EnzymeCore.Annotation{<:FunctionWrappersWrapper}, @@ -338,9 +351,8 @@ function EnzymeRules.reverse( args::Vararg{EnzymeCore.Annotation, N} ) where {N} f_orig = unwrap(func.val) - # Only worth invoking Enzyme.autodiff when at least one arg is - # Duplicated/BatchDuplicated — otherwise there's nothing to accumulate. - if any(a -> a isa EnzymeCore.Duplicated || a isa EnzymeCore.BatchDuplicated, args) + # Nothing to do if every argument is Const. + if any(a -> !(a isa EnzymeCore.Const), args) # Temporarily restore the call-time argument values (snapshotted in the # tape during augmented_primal) so the VJP is taken about the correct # state, then put the live values back so the surrounding reverse pass @@ -349,16 +361,13 @@ function EnzymeRules.reverse( # recomputation differentiate `f_orig` about the wrong state. live = ntuple(i -> (tape[i] === nothing ? nothing : _snapshot(args[i].val)), Val(N)) map((a, snap) -> _restore!(a.val, snap), args, tape) - Enzyme.autodiff(Reverse, Const(f_orig), Const, args...) + raw = Enzyme.autodiff(Reverse, Const(f_orig), Const, args...)[1]::NTuple{N, Any} map((a, snap) -> _restore!(a.val, snap), args, live) + # `map` over tuples specialises per element, so dispatching `_revslot` + # on each arg's concrete annotation type yields an exactly-typed result. + return map(_revslot, args, raw) end - return ntuple(Val(N)) do i - if args[i] isa EnzymeCore.Active - zero(eltype(typeof(args[i]))) - else - nothing - end - end + return ntuple(_ -> nothing, Val(N)) end end diff --git a/test/Enzyme/enzyme_tests.jl b/test/Enzyme/enzyme_tests.jl index b08d25b..9696c0b 100644 --- a/test/Enzyme/enzyme_tests.jl +++ b/test/Enzyme/enzyme_tests.jl @@ -402,6 +402,33 @@ end @test g ≈ fd rtol = 1.0e-4 end +@testset "Enzyme Reverse: IIP wrapper with a mix of Duplicated and Active args" begin + # A time-dependent in-place rhs, differentiated so the reverse rule sees + # (Duplicated du, Duplicated u, Duplicated p, Active t). The rule must + # return the *real* gradient for the Active `t` with an exact-typed tuple + # (Nothing per Duplicated arg, Float64 for the Active). Before the fix this + # returned a union-typed `(nothing, …, 0.0)` — Enzyme rejected it with a + # `ReverseRuleReturnError`, and the `t`-gradient was zeroed rather than + # computed. + f!(du, u, p, t) = (du[1] = p[1] * u[1] + t * u[2]; du[2] = p[2] * u[2]; nothing) + ARGT = Tuple{Vector{Float64}, Vector{Float64}, Vector{Float64}, Float64} + + function loss(x) # x = [p1, p2, t] + u = [1.5, 2.0] + du = zero(u) + wf = FunctionWrappersWrapper(f!, (ARGT,), (Nothing,)) + wf(du, u, [x[1], x[2]], x[3]) # t = x[3] flows in as an Active scalar + return du[1]^2 + du[2]^2 + end + + x = [0.7, 0.4, 0.9] + g = collect(Enzyme.gradient(Enzyme.set_runtime_activity(Enzyme.Reverse), loss, x)[1]) + # du = [p1*u1 + t*u2, p2*u2]; loss = du1^2 + du2^2 + du1 = x[1] * 1.5 + x[3] * 2.0 + du2 = x[2] * 2.0 + @test g ≈ [2 * du1 * 1.5, 2 * du2 * 2.0, 2 * du1 * 2.0] # ∂/∂t = 2*du1*u2 ≠ 0 +end + # ============================================================================= # Runtime-activity propagation through the FWW forward rules. #