Merge pull request #1348 from ChrisRackauckas-Claude/normalize-callback-vjp-args

Normalize CallbackAffectPWrapper VJP argument ordering

Author: ChrisRackauckas

src/callback_tracking.jl

@@ -382,15 +382,15 @@ function _setup_reverse_callbacks(
 
                 if sensealg isa GaussAdjoint
                     vecjacobian!(
-                        dgrad, integrator.p,
+                        nothing, y,
                         integrator.f.f.integrating_cb.affect!.integrand_values.integrand,
-                        y, integrator.t, fakeSp; dgrad = nothing, dy = nothing
+                        integrator.p, integrator.t, fakeSp; dgrad = dgrad, dy = nothing
                     )
                     integrator.f.f.integrating_cb.affect!.integrand_values.integrand .= dgrad
                 else
                     vecjacobian!(
-                        dgrad, integrator.p, grad, y, integrator.t, fakeSp;
-                        dgrad = nothing, dy = nothing
+                        nothing, y, grad, integrator.p, integrator.t, fakeSp;
+                        dgrad = dgrad, dy = nothing
                     )
                     grad .= dgrad
                 end
@@ -483,7 +483,7 @@ mutable struct CallbackAffectPWrapper{cbType, AJV, EI, T} <: Function
     tprev::T
 end
 
-function (ff::CallbackAffectPWrapper)(dp, p, u, t)
+function (ff::CallbackAffectPWrapper)(dp, u, p, t)
     _affect! = get_affect!(ff.cb, ff.pos_neg)
     fakeinteg = get_FakeIntegrator(ff.autojacvec, u, p, t, ff.tprev)
     if ff.cb isa VectorContinuousCallback
@@ -510,6 +510,30 @@ function get_FakeIntegrator(autojacvec::ReverseDiffVJP, u, p, t, tprev)
 end
 get_FakeIntegrator(autojacvec::EnzymeVJP, 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))
+end
+
+function _get_wp_paramjac_config(autojacvec::ReverseDiffVJP, _p, wp, y, __p, _t)
+    if _p === nothing || _p isa SciMLBase.NullParameters
+        tunables, repack = _p, identity
+    else
+        tunables, repack, aliases = canonicalize(Tunable(), _p)
+    end
+    tunables_inner = tunables
+    tape = ReverseDiff.GradientTape((y, tunables_inner, [_t])) do u, p, t
+        dp1 = similar(p, length(p))
+        dp1 .= false
+        wp(dp1, u, repack(p), first(t))
+        return vec(dp1)
+    end
+    if compile_tape(autojacvec)
+        return ReverseDiff.compile(tape)
+    else
+        return tape
+    end
+end
+
 function get_cb_diffcaches(
         cb::Union{
             DiscreteCallback, ContinuousCallback,
@@ -562,11 +586,8 @@ function get_cb_diffcaches(
                     nothing, nothing, nothing, false
                 )
 
-                paramjac_config = get_paramjac_config(
-                    autojacvec, y, wp, _p, y, _t;
-                    numindvar = length(y), alg = nothing,
-                    isinplace = true, isRODE = false,
-                    _W = nothing
+                paramjac_config = _get_wp_paramjac_config(
+                    autojacvec, _p, wp, y, _p, _t
                 )
                 pf = get_pf(autojacvec; _f = wp, isinplace = true, isRODE = false)
                 if autojacvec isa EnzymeVJP

src/derivative_wrappers.jl

@@ -492,8 +492,10 @@ function _vecjacobian!(
         _y = eltype(y) === eltype(λ) ? y : convert.(promote_type(eltype(y), eltype(λ)), y)
         if W === nothing
             _tunables, _repack, _ = canonicalize(Tunable(), _p)
+            _is_pswap = TS <: CallbackSensitivityFunctionPSwap
             tape = ReverseDiff.GradientTape((_y, _tunables, [t])) do u, p, t
-                du1 = similar(u, size(u))
+                du1 = _is_pswap ? similar(p, length(p)) : similar(u, size(u))
+                du1 .= false
                 f(du1, u, _repack(p), first(t))
                 return vec(du1)
             end

src/quadrature_adjoint.jl

@@ -548,14 +548,11 @@ function _update_integrand_and_dgrad(
 
     _p = similar(integrand.p, size(integrand.p))
     _p .= false
-    wp(_p, integrand.p, integrand.y, t)
+    wp(_p, integrand.y, integrand.p, t)
 
     if _p != integrand.p
-        paramjac_config = get_paramjac_config(
-            sensealg.autojacvec, integrand.y, wp, _p, integrand.y, t;
-            numindvar = length(integrand.y), alg = nothing,
-            isinplace = true, isRODE = false,
-            _W = nothing
+        paramjac_config = _get_wp_paramjac_config(
+            sensealg.autojacvec, integrand.p, wp, integrand.y, integrand.p, t
         )
         pf = get_pf(sensealg.autojacvec; _f = wp, isinplace = true, isRODE = false)
         if sensealg.autojacvec isa EnzymeVJP
@@ -572,8 +569,8 @@ function _update_integrand_and_dgrad(
         fakeSp = CallbackSensitivityFunctionPSwap(wp, sensealg, diffcache_wp, sol.prob)
         #vjp with Jacobin given by dw/dp before event and vector given by grad
         vecjacobian!(
-            res, integrand.p, res, integrand.y, t, fakeSp;
-            dgrad = nothing, dy = nothing
+            nothing, integrand.y, res, integrand.p, t, fakeSp;
+            dgrad = res, dy = nothing
         )
         integrand = update_p_integrand(integrand, _p)
     end