Merge pull request #1222 from SciML/void

Use SciMLBase.Void on EnzymeVJP functions

Author: ChrisRackauckas

src/adjoint_common.jl

@@ -489,7 +489,7 @@ function get_pf(autojacvec::ReverseDiffVJP; _f = nothing, isinplace = nothing,
 end
 
 function get_pf(autojacvec::EnzymeVJP; _f, isinplace, isRODE)
-    _f
+    isinplace ? SciMLBase.Void(_f) : _f
 end
 
 function get_pf(::MooncakeVJP, prob, _f)

src/derivative_wrappers.jl

@@ -724,19 +724,19 @@ function _vecjacobian!(dλ, y, λ, p, t, S::TS, isautojacvec::EnzymeVJP, dgrad,
             # 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)
+            _tmp6 = Enzyme.make_zero(SciMLBase.Void(f))
         else
             Enzyme.remake_zero!(_tmp6)
         end
 
         if W === nothing
-            Enzyme.autodiff(Enzyme.Reverse, Enzyme.Duplicated(f, _tmp6),
+            Enzyme.autodiff(Enzyme.Reverse, Enzyme.Duplicated(SciMLBase.Void(f), _tmp6),
                 Enzyme.Const, Enzyme.Duplicated(tmp3, tmp4),
                 Enzyme.Duplicated(ytmp, tmp1),
                 dup,
                 Enzyme.Const(t))
         else
-            Enzyme.autodiff(Enzyme.Reverse, Enzyme.Duplicated(f, _tmp6),
+            Enzyme.autodiff(Enzyme.Reverse, Enzyme.Duplicated(SciMLBase.Void(f), _tmp6),
                 Enzyme.Const, Enzyme.Duplicated(tmp3, tmp4),
                 Enzyme.Duplicated(ytmp, tmp1),
                 dup,

src/gauss_adjoint.jl

@@ -415,7 +415,7 @@ function GaussIntegrand(sol, sensealg, checkpoints, dgdp = nothing)
         pf = nothing
         pJ = nothing
     elseif sensealg.autojacvec isa EnzymeVJP
-        pf = unwrappedf
+        pf = SciMLBase.isinplace(sol.prob.f) ? SciMLBase.Void(unwrappedf) : unwrappedf
         paramjac_config = zero(y), zero(y), Enzyme.make_zero(pf)
         pJ = nothing
     elseif sensealg.autojacvec isa MooncakeVJP
@@ -616,4 +616,4 @@ function _adjoint_sensitivities(sol, sensealg::GaussAdjoint, alg; t = nothing,
 end
 
 __maybe_adjoint(x::AbstractArray) = x'
-__maybe_adjoint(x) = x
+__maybe_adjoint(x) = x

src/quadrature_adjoint.jl

@@ -209,7 +209,7 @@ function AdjointSensitivityIntegrand(sol, adj_sol, sensealg, dgdp = nothing)
         pf = nothing
         pJ = nothing
     elseif sensealg.autojacvec isa EnzymeVJP
-        pf = unwrappedf
+        pf = SciMLBase.isinplace(sol.prob.f) ? SciMLBase.Void(unwrappedf) : unwrappedf
         paramjac_config = zero(y), zero(y), Enzyme.make_zero(pf)
         pJ = nothing
     elseif sensealg.autojacvec isa MooncakeVJP
@@ -293,7 +293,7 @@ function vec_pjac!(out, λ, y, t, S::AdjointSensitivityIntegrand)
         if SciMLBase.isinplace(sol.prob.f)
             Enzyme.remake_zero!(tmp6)
             Enzyme.autodiff(
-                Enzyme.Reverse, Enzyme.Duplicated(f, tmp6), Enzyme.Const,
+                Enzyme.Reverse, Enzyme.Duplicated(SciMLBase.Void(f), tmp6), Enzyme.Const,
                 Enzyme.Duplicated(tmp3, tmp4),
                 Enzyme.Const(y), dup, Enzyme.Const(t))
         else
@@ -537,4 +537,4 @@ function _update_integrand_and_dgrad(res, sensealg::QuadratureAdjoint, cb, integ
     vecjacobian!(dλ, integrand.y, dλ, integrand.p, t, fakeS; dgrad = dgrad)
     res .-= dgrad
     return integrand
-end
+end