Safety make_zero! on repeated Enzyme calls with caches

This should ensure that the caches are always zero'd for the derivatives.
SciML · Jun 28, 2024 · 0c56bec · 0c56bec
1 parent 456ea4b
commit 0c56bec
Show file tree

Hide file tree

Showing 4 changed files with 5 additions and 1 deletion.
diff --git a/Project.toml b/Project.toml
@@ -61,7 +61,7 @@ DiffEqNoiseProcess = "5.19"
 Distributed = "1"
 Distributions = "0.25"
 EllipsisNotation = "1"
-Enzyme = "0.12"
+Enzyme = "0.12.12"
 FiniteDiff = "2"
 ForwardDiff = "0.10"
 FunctionProperties = "0.1"

diff --git a/src/derivative_wrappers.jl b/src/derivative_wrappers.jl
@@ -705,6 +705,8 @@ function _vecjacobian!(dλ, y, λ, p, t, S::TS, isautojacvec::EnzymeVJP, dgrad,
 
     isautojacvec = get_jacvec(sensealg)
 
+    Enzyme.make_zero!(_tmp6)
+
     if inplace_sensitivity(S)
         if W === nothing
             Enzyme.autodiff(Enzyme.Reverse, Enzyme.Duplicated(S.diffcache.pf, _tmp6),

diff --git a/src/gauss_adjoint.jl b/src/gauss_adjoint.jl
@@ -498,6 +498,7 @@ function vec_pjac!(out, λ, y, t, S::GaussIntegrand)
         vtmp4 = vec(tmp4)
         vtmp4 .= λ
         out .= 0
+        Enzyme.make_zero!(tmp6)
         Enzyme.autodiff(
             Enzyme.Reverse, Enzyme.Duplicated(pf, tmp6), Enzyme.Const,
             Enzyme.Duplicated(tmp3, tmp4),

diff --git a/src/quadrature_adjoint.jl b/src/quadrature_adjoint.jl
@@ -294,6 +294,7 @@ function vec_pjac!(out, λ, y, t, S::AdjointSensitivityIntegrand)
         tmp3, tmp4, tmp6 = paramjac_config
         tmp4 .= λ
         out .= 0
+        Enzyme.make_zero!(tmp6)
         Enzyme.autodiff(
             Enzyme.Reverse, Enzyme.Duplicated(pf, tmp6), Enzyme.Const,
             Enzyme.Duplicated(tmp3, tmp4),