Merge pull request #1191 from jClugstor/enzyme_adjoint_lbc

use LazyBufferCache instead of FixedSizeDiffCache

Author: ChrisRackauckas

src/SciMLSensitivity.jl

@@ -16,7 +16,7 @@ using FunctionWrappersWrappers: FunctionWrappersWrappers
 using GPUArraysCore: GPUArraysCore
 using LinearSolve: LinearSolve
 using PreallocationTools: PreallocationTools, dualcache, get_tmp, DiffCache,
-                          FixedSizeDiffCache
+                          FixedSizeDiffCache, LazyBufferCache
 using RandomNumbers: Xorshifts
 using RecursiveArrayTools: RecursiveArrayTools, AbstractDiffEqArray,
                            AbstractVectorOfArray, ArrayPartition, DiffEqArray,

src/adjoint_common.jl

@@ -444,10 +444,10 @@ function get_paramjac_config(autojacvec::EnzymeVJP, p::SciMLBase.NullParameters,
             autojacvec.chunksize
         end
 
-        paramjac_config = FixedSizeDiffCache(zero(y), chunk), p,
-        FixedSizeDiffCache(zero(y), chunk),
-        FixedSizeDiffCache(zero(y), chunk),
-        FixedSizeDiffCache(zero(y), chunk)
+        paramjac_config = LazyBufferCache(), p, 
+        LazyBufferCache(),
+        LazyBufferCache(),
+        LazyBufferCache()
     else
         paramjac_config = zero(y), p, zero(y), zero(y), zero(y)
     end
@@ -464,11 +464,10 @@ function get_paramjac_config(autojacvec::EnzymeVJP, p, f, y, _p, _t; numindvar,
             autojacvec.chunksize
         end
 
-        paramjac_config = FixedSizeDiffCache(zero(y), chunk),
-        zero(_p),
-        FixedSizeDiffCache(zero(y), chunk),
-        FixedSizeDiffCache(zero(y), chunk),
-        FixedSizeDiffCache(zero(y), chunk)
+        paramjac_config = LazyBufferCache(), zero(_p), 
+        LazyBufferCache(),
+        LazyBufferCache(),
+        LazyBufferCache()
     else
         paramjac_config = zero(y), zero(_p), zero(y), zero(y), zero(y)
     end

src/derivative_wrappers.jl

@@ -677,7 +677,7 @@ function _vecjacobian!(dλ, y, λ, p, t, S::TS, isautojacvec::EnzymeVJP, dgrad,
 
     _tmp1, tmp2, _tmp3, _tmp4, _tmp5, _tmp6 = S.diffcache.paramjac_config
 
-    if _tmp1 isa FixedSizeDiffCache
+    if _tmp1 isa LazyBufferCache
         tmp1 = get_tmp(_tmp1, dλ)
         tmp3 = get_tmp(_tmp3, dλ)
         tmp4 = get_tmp(_tmp4, dλ)