Use ConstantOrCache for parameters in JVPCache

Co-authored-by: Guillaume Dalle <[email protected]>

Author: jClugstor

lib/OrdinaryDiffEqDifferentiation/src/operators.jl

@@ -66,10 +66,8 @@ function prepare_jvp(f::DiffEqBase.AbstractDiffEqFunction, du, u, p, t, autodiff
     SciMLBase.has_jvp(f) && return f.jvp
     autodiff = autodiff isa AutoSparse ?  ADTypes.dense_ad(autodiff) : autodiff
     @assert DI.check_inplace(autodiff) "AD backend $(autodiff) doesn't support in-place problems."
-    di_prep = DI.prepare_pushforward(
-        (du, u) -> f(du,u,p,t), du, autodiff, u, (u,))
-    return (Jv, v, u, p, t) -> DI.pushforward!((du,x) -> f(du,x,p,t), du, (reshape(Jv, size(du)),), di_prep,
-            autodiff, u, (reshape(v,size(u)),))
+    di_prep = DI.prepare_pushforward(f, du, autodiff, u, (u,), DI.ConstantOrCache(p), DI.Constant(t))
+    return (Jv, v, u, p, t) -> DI.pushforward!(f, du, (reshape(Jv, size(du)),), di_prep, autodiff, u, (reshape(v,size(u)), DI.ConstantOrCache(p), DI.Constant(t)))
 end
 
 function SciMLOperators.update_coefficients!(J::JVPCache, u, p, t)