fix function wrap

Author: ChrisRackauckas

src/adjoint_common.jl

@@ -489,30 +489,7 @@ function get_pf(autojacvec::ReverseDiffVJP; _f = nothing, isinplace = nothing,
 end
 
 function get_pf(autojacvec::EnzymeVJP; _f, isinplace, isRODE)
-    pf = let f = _f
-        if isinplace && isRODE
-            function (out, u, _p, t, W)
-                f(out, u, _p, t, W)
-                nothing
-            end
-        elseif isinplace
-            function (out, u, _p, t)
-                f(out, u, _p, t)
-                nothing
-            end
-        elseif !isinplace && isRODE
-            function (out, u, _p, t, W)
-                out .= f(u, _p, t, W)
-                nothing
-            end
-        else
-            # !isinplace
-            function (out, u, _p, t)
-                out .= f(u, _p, t)
-                nothing
-            end
-        end
-    end
+    _f
 end
 
 function get_pf(::MooncakeVJP, prob, _f)

src/gauss_adjoint.jl

@@ -16,7 +16,7 @@ end
 struct ODEGaussAdjointSensitivityFunction{C <: AdjointDiffCache,
     Alg <: GaussAdjoint,
     uType, SType, CPS, pType,
-    fType <: DiffEqBase.AbstractDiffEqFunction,
+    fType,
     GI <: GaussIntegrand,
     ICB} <: SensitivityFunction
     diffcache::C
@@ -82,7 +82,7 @@ function ODEGaussAdjointSensitivityFunction(
         nothing
     end
     diffcache, y = adjointdiffcache(
-        g, sensealg, discrete, sol, dgdu, dgdp, sol.prob.f, alg;
+        g, sensealg, discrete, sol, dgdu, dgdp, f, alg;
         quad = true)
     return ODEGaussAdjointSensitivityFunction(diffcache, sensealg, discrete,
         y, sol, checkpoint_sol, sol.prob, f, gaussint, integrating_cb)
@@ -415,20 +415,7 @@ function GaussIntegrand(sol, sensealg, checkpoints, dgdp = nothing)
         pf = nothing
         pJ = nothing
     elseif sensealg.autojacvec isa EnzymeVJP
-        pf = let f = unwrappedf
-            if DiffEqBase.isinplace(prob)
-                function (out, u, _p, t)
-                    f(out, u, _p, t)
-                    nothing
-                end
-            else
-                !DiffEqBase.isinplace(prob)
-                function (out, u, _p, t)
-                    out .= f(u, _p, t)
-                    nothing
-                end
-            end
-        end
+        pf = unwrappedf
         paramjac_config = zero(y), zero(y), Enzyme.make_zero(pf)
         pJ = nothing
     elseif sensealg.autojacvec isa MooncakeVJP

src/quadrature_adjoint.jl

@@ -209,30 +209,7 @@ function AdjointSensitivityIntegrand(sol, adj_sol, sensealg, dgdp = nothing)
         pf = nothing
         pJ = nothing
     elseif sensealg.autojacvec isa EnzymeVJP
-        pf = let f = unwrappedf
-            if DiffEqBase.isinplace(prob) && prob isa RODEProblem
-                function (out, u, _p, t, W)
-                    f(out, u, _p, t, W)
-                    nothing
-                end
-            elseif DiffEqBase.isinplace(prob)
-                function (out, u, _p, t)
-                    f(out, u, _p, t)
-                    nothing
-                end
-            elseif !DiffEqBase.isinplace(prob) && prob isa RODEProblem
-                function (out, u, _p, t, W)
-                    out .= f(u, _p, t, W)
-                    nothing
-                end
-            else
-                !DiffEqBase.isinplace(prob)
-                function (out, u, _p, t)
-                    out .= f(u, _p, t)
-                    nothing
-                end
-            end
-        end
+        pf = unwrappedf
         paramjac_config = zero(y), zero(y), Enzyme.make_zero(pf)
         pJ = nothing
     elseif sensealg.autojacvec isa MooncakeVJP