Merge pull request #1221 from SciML/remake_zero2

actually add remake_zero

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/derivative_wrappers.jl

@@ -720,12 +720,15 @@ function _vecjacobian!(dλ, y, λ, p, t, S::TS, isautojacvec::EnzymeVJP, dgrad,
     isautojacvec = get_jacvec(sensealg)
 
     if inplace_sensitivity(S)
-
-        # Correctness over speed
-        # TODO: Get a fix for `remake_zero!` to allow reusing zero'd memory
-        # https://github.com/EnzymeAD/Enzyme.jl/issues/2400
-        _tmp6 = Enzyme.make_zero(f)
-
+        if S isa CallbackSensitivityFunction
+            # Correctness over speed
+            # TODO: Get a fix for `remake_zero!` to allow reusing zero'd memory
+            # https://github.com/EnzymeAD/Enzyme.jl/issues/2400
+            _tmp6 = Enzyme.make_zero(f)
+        else
+            Enzyme.remake_zero!(_tmp6)
+        end
+        
         if W === nothing
             Enzyme.autodiff(Enzyme.Reverse, Enzyme.Duplicated(f, _tmp6),
                 Enzyme.Const, Enzyme.Duplicated(tmp3, tmp4),

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
@@ -508,10 +495,7 @@ function vec_pjac!(out, λ, y, t, S::GaussIntegrand)
         Enzyme.remake_zero!(out)
 
         if SciMLBase.isinplace(sol.prob.f)
-            # Correctness over speed
-            # TODO: Get a fix for `remake_zero!` to allow reusing zero'd memory
-            # https://github.com/EnzymeAD/Enzyme.jl/issues/2400
-            tmp6 = Enzyme.make_zero(tmp6)
+            Enzyme.remake_zero!(tmp6)
 
             Enzyme.autodiff(
                 Enzyme.Reverse, Enzyme.Duplicated(pf, tmp6), Enzyme.Const,

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
@@ -314,10 +291,7 @@ function vec_pjac!(out, λ, y, t, S::AdjointSensitivityIntegrand)
         end
 
         if SciMLBase.isinplace(sol.prob.f)
-            # Correctness over speed
-            # TODO: Get a fix for `remake_zero!` to allow reusing zero'd memory
-            # https://github.com/EnzymeAD/Enzyme.jl/issues/2400
-            tmp6 = Enzyme.make_zero(f)
+            Enzyme.remake_zero!(tmp6)
             Enzyme.autodiff(
                 Enzyme.Reverse, Enzyme.Duplicated(f, tmp6), Enzyme.Const,
                 Enzyme.Duplicated(tmp3, tmp4),