Merge pull request #1421 from ChrisRackauckas-Claude/mooncakevjp-callback-compat

Add MooncakeVJP to callback-compatible VJP methods

Author: ChrisRackauckas

Project.toml

@@ -86,7 +86,7 @@ Lux = "1"
 Markdown = "1.10"
 ModelingToolkit = "10, 11"
 ModelingToolkitStandardLibrary = "2"
-Mooncake = "0.5"
+Mooncake = "0.5.25"
 Reactant = "0.2.22"
 NLsolve = "4.5.1"
 NonlinearSolve = "3.0.1, 4"

ext/SciMLSensitivityMooncakeExt.jl

@@ -1,8 +1,9 @@
 module SciMLSensitivityMooncakeExt
 
-using SciMLSensitivity: SciMLSensitivity
+using SciMLSensitivity: SciMLSensitivity, FakeIntegrator
 using Mooncake: Mooncake
-import SciMLSensitivity: get_paramjac_config, mooncake_run_ad, MooncakeVJP, MooncakeLoaded,
+import SciMLSensitivity: get_paramjac_config, get_cb_paramjac_config, mooncake_run_ad,
+    MooncakeVJP, MooncakeLoaded,
     DiffEqBase, MooncakeAdjoint, _init_originator_gradient
 using SciMLSensitivity: SciMLBase, SciMLStructures, canonicalize, Tunable, isscimlstructure,
     SciMLStructuresCompatibilityError, convert_tspan,
@@ -35,6 +36,68 @@ function get_paramjac_config(::MooncakeLoaded, ::MooncakeVJP, pf, p, f, y, _t)
     return cache, pf, λ_mem, dy_mem, p_grad_buf
 end
 
+"""
+    get_cb_paramjac_config(::MooncakeLoaded, ::MooncakeVJP, raw_affect, event_idx, y, p, _t, mode)
+
+Build a Mooncake pullback cache for a tracked callback affect function. Mirrors
+the `get_cb_paramjac_config(::ReactantLoaded, ::ReactantVJP, ...)` entry point:
+`raw_affect` is extracted upfront (`get_affect!(cb, pos_neg)` at the call site)
+so the Mooncake-traced closure does not need to recursively unwrap
+`TrackedAffect`, which would otherwise trip on the `Base.argument_datatype`
+ccall surfaced by that dispatch.
+
+`mode === :state` builds a cache for the state-affect closure (state-sized
+output); `mode === :param` builds one for the parameter-affect closure
+(parameter-sized output) so its Mooncake cotangent/output buffers match the
+flat tunables shape rather than the state shape. The returned 5-tuple has the
+same layout as `get_paramjac_config(::MooncakeLoaded, ::MooncakeVJP, ...)` so
+`_vecjacobian!(::MooncakeVJP)` / `mooncake_run_ad` can consume it unchanged.
+"""
+function get_cb_paramjac_config(
+        ::MooncakeLoaded, ::MooncakeVJP, raw_affect, event_idx, y, p, _t, mode
+    )
+    has_event_idx = event_idx !== nothing
+    tprev0 = _t
+
+    if mode === :state
+        pf = let raw = raw_affect, ev = event_idx, tprev = tprev0, has_ev = has_event_idx
+            (out, u, p, t) -> begin
+                fakeinteg = FakeIntegrator(copy(u), copy(p), t, tprev)
+                if has_ev
+                    raw(fakeinteg, ev)
+                else
+                    raw(fakeinteg)
+                end
+                copyto!(out, fakeinteg.u)
+                return out
+            end
+        end
+        out_sample = y
+    elseif mode === :param
+        pf = let raw = raw_affect, ev = event_idx, tprev = tprev0, has_ev = has_event_idx
+            (out, u, p, t) -> begin
+                fakeinteg = FakeIntegrator(copy(u), copy(p), t, tprev)
+                if has_ev
+                    raw(fakeinteg, ev)
+                else
+                    raw(fakeinteg)
+                end
+                copyto!(out, fakeinteg.p)
+                return out
+            end
+        end
+        out_sample = p
+    else
+        error("get_cb_paramjac_config: unknown mode $(mode); expected :state or :param")
+    end
+
+    dy_mem = zero(out_sample)
+    λ_mem = zero(out_sample)
+    cache = Mooncake.prepare_pullback_cache(pf, dy_mem, y, p, _t)
+    p_grad_buf = p isa AbstractArray && !(p isa Array) ? similar(p) : nothing
+    return cache, pf, λ_mem, dy_mem, p_grad_buf
+end
+
 function mooncake_run_ad(paramjac_config::Tuple, y, p, t, λ)
     cache, pf, λ_mem, dy_mem, p_grad_buf = paramjac_config
     λ_mem .= λ

src/adjoint_common.jl

@@ -666,6 +666,12 @@ function get_cb_paramjac_config(::Any, ::ReactantVJP, raw_affect, event_idx, y,
     error(msg)
 end
 
+function get_cb_paramjac_config(::Any, ::MooncakeVJP, raw_affect, event_idx, y, p, _t, mode)
+    msg = "MooncakeVJP requires Mooncake.jl is loaded. Install the package and do " *
+        "`using Mooncake` to use this functionality"
+    error(msg)
+end
+
 function getprob(S::SensitivityFunction)
     return (S isa ODEBacksolveSensitivityFunction) ? S.prob : S.sol.prob
 end

src/callback_tracking.jl

@@ -283,11 +283,12 @@ function _setup_reverse_callbacks(
     # if save_positions = [1,0] the gradient contribution is added before, and in principle we would need to correct the adjoint state again. Therefore,
 
     cb.save_positions == [1, 0] && error("save_positions=[1,0] is currently not supported.")
-    # Callbacks require ReverseDiffVJP or EnzymeVJP for their own VJP computations.
-    # The ODE adjoint may use a different autojacvec (even numerical/false), but the
-    # callback affect functions (CallbackAffectWrapper) are separate and typically work
-    # with ReverseDiff even when the ODE function doesn't.
-    cb_autojacvec = if sensealg.autojacvec isa Union{ReverseDiffVJP, EnzymeVJP, ReactantVJP}
+    # Callback affect functions (CallbackAffectWrapper) are traced separately
+    # from the ODE function, so the ODE adjoint may use a different autojacvec
+    # (even numerical/false) while the callback path uses its own compatible
+    # backend. The `supports_callback_vjp` trait marks which VJP backends have
+    # a dedicated callback path; anything else falls back to ReverseDiffVJP.
+    cb_autojacvec = if supports_callback_vjp(sensealg.autojacvec)
         sensealg.autojacvec
     else
         @warn "autojacvec=$(sensealg.autojacvec) is not compatible with callbacks, using ReverseDiffVJP() for callback VJPs"
@@ -506,9 +507,10 @@ end
 
 function setup_w_wp(
         cb::Union{DiscreteCallback, ContinuousCallback, VectorContinuousCallback},
-        autojacvec::Union{ReverseDiffVJP, EnzymeVJP, ReactantVJP}, pos_neg, event_idx,
-        tprev
+        autojacvec, pos_neg, event_idx, tprev
     )
+    supports_callback_vjp(autojacvec) ||
+        error("setup_w_wp called with a VJP backend that does not support callbacks: $(autojacvec). This is an internal error — the callback path should have redirected to a compatible backend in `_setup_reverse_callbacks`.")
     w = CallbackAffectWrapper(cb, autojacvec, pos_neg, event_idx, tprev)
     wp = CallbackAffectPWrapper(cb, autojacvec, pos_neg, event_idx, tprev)
     return w, wp
@@ -519,6 +521,7 @@ function get_FakeIntegrator(autojacvec::ReverseDiffVJP, u, p, t, tprev)
 end
 get_FakeIntegrator(autojacvec::EnzymeVJP, u, p, t, tprev) = FakeIntegrator(u, p, t, tprev)
 get_FakeIntegrator(autojacvec::ReactantVJP, u, p, t, tprev) = FakeIntegrator(u, p, t, tprev)
+get_FakeIntegrator(autojacvec::MooncakeVJP, u, p, t, tprev) = FakeIntegrator(u, p, t, tprev)
 
 function _get_wp_paramjac_config(autojacvec::EnzymeVJP, _p, wp, y, __p, _t)
     return (zero(y), zero(_p), zero(_p), zero(_p), zero(y))
@@ -603,6 +606,41 @@ function get_cb_diffcaches(
                         nothing, nothing, nothing, false,
                         nothing, identity
                     )
+                elseif autojacvec isa MooncakeVJP
+                    # MooncakeVJP: build Mooncake pullback caches for the
+                    # state-affect and parameter-affect callback wrappers via
+                    # `get_cb_paramjac_config` in the Mooncake extension, mirroring
+                    # the ReactantVJP branch above. Mooncake can't trace through
+                    # the recursive TrackedAffect unwrapping in
+                    # `CallbackAffectWrapper` (it trips on the
+                    # `Base.argument_datatype` ccall surfaced by that dispatch),
+                    # so the extension hoists `raw_affect` out and bakes it into
+                    # a flat (out, u, p, t) closure before preparing the cache.
+                    raw_affect = get_affect!(cb, pos_neg)
+
+                    w_paramjac = get_cb_paramjac_config(
+                        MooncakeLoaded(), autojacvec, raw_affect,
+                        event_idx, y, _p, _t, :state
+                    )
+                    diffcache_w = AdjointDiffCache(
+                        nothing, nothing, nothing, nothing, nothing,
+                        nothing, nothing, nothing, w_paramjac,
+                        nothing, nothing, nothing, nothing, nothing,
+                        nothing, nothing, nothing, false,
+                        nothing, identity
+                    )
+
+                    wp_paramjac = get_cb_paramjac_config(
+                        MooncakeLoaded(), autojacvec, raw_affect,
+                        event_idx, y, _p, _t, :param
+                    )
+                    diffcache_wp = AdjointDiffCache(
+                        nothing, nothing, nothing, nothing, nothing,
+                        nothing, nothing, nothing, wp_paramjac,
+                        nothing, nothing, nothing, nothing, nothing,
+                        nothing, nothing, nothing, false,
+                        nothing, identity
+                    )
                 else
                     w, wp = setup_w_wp(cb, autojacvec, pos_neg, event_idx, _t)
 

src/sensitivity_algorithms.jl

@@ -1696,6 +1696,23 @@ supports_structured_vjp(::ReverseDiffVJP) = false
 supports_structured_vjp(::Bool) = false
 supports_structured_vjp(::Nothing) = false
 
+"""
+    supports_callback_vjp(autojacvec) -> Bool
+
+Return `true` if the VJP backend can differentiate the affect functions of
+tracked callbacks (i.e. has a dedicated branch in `get_cb_diffcaches` and a
+`_vecjacobian!` path that works with `CallbackSensitivityFunction`).
+
+When `false`, the callback adjoint code falls back to `ReverseDiffVJP()` for
+its own VJP computations — the ODE adjoint may still use the user-requested
+backend, since the callback affect functions are traced separately.
+"""
+supports_callback_vjp(::ReverseDiffVJP) = true
+supports_callback_vjp(::EnzymeVJP) = true
+supports_callback_vjp(::ReactantVJP) = true
+supports_callback_vjp(::MooncakeVJP) = true
+supports_callback_vjp(::Any) = false
+
 """
 ```julia
 ForwardDiffOverAdjoint{A} <: AbstractSecondOrderSensitivityAlgorithm{nothing, true, nothing}

test/callbacks/continuous_callbacks.jl

@@ -1,5 +1,6 @@
-using OrdinaryDiffEq, Zygote, Reactant
+using OrdinaryDiffEq, Zygote, Reactant, Mooncake
 using SciMLSensitivity, Test, ForwardDiff, FiniteDiff
+using SciMLSensitivity: MooncakeVJP
 
 abstol = 1.0e-12
 reltol = 1.0e-12
@@ -362,5 +363,14 @@ println("Continuous Callbacks")
         sensealg = GaussAdjoint(autojacvec = ReactantVJP(allow_scalar = true))
         gZy = Zygote.gradient(p -> loss(p, cb, sensealg), p)[1]
         @test gFD ≈ gZy rtol = 1.0e-10
+
+        # MooncakeVJP with callbacks
+        sensealg = InterpolatingAdjoint(autojacvec = MooncakeVJP())
+        gZy = Zygote.gradient(p -> loss(p, cb, sensealg), p)[1]
+        @test gFD ≈ gZy rtol = 1.0e-10
+
+        sensealg = GaussAdjoint(autojacvec = MooncakeVJP())
+        gZy = Zygote.gradient(p -> loss(p, cb, sensealg), p)[1]
+        @test gFD ≈ gZy rtol = 1.0e-10
     end
 end