Resolve JET.jl issue

nathanaelbosch · Aug 30, 2024 · a5030a2 · a5030a2
1 parent 7ff97e7
commit a5030a2
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/src/caches.jl b/src/caches.jl
@@ -251,7 +251,9 @@ function OrdinaryDiffEq.alg_cache(
     )
 end
 
-get_uf(f, t, p, ::Val{true}) = OrdinaryDiffEq.UJacobianWrapper(f, t, p)
-get_uf(f, t, p, ::Val{false}) = OrdinaryDiffEq.UDerivativeWrapper(f, t, p)
+get_uf(f, t, p, ::Val{true}) =
+    OrdinaryDiffEq.OrdinaryDiffEqDifferentiation.UJacobianWrapper(f, t, p)
+get_uf(f, t, p, ::Val{false}) =
+    OrdinaryDiffEq.OrdinaryDiffEqDifferentiation.UDerivativeWrapper(f, t, p)
 
 OrdinaryDiffEq.get_fsalfirstlast(cache::EKCache, rate_prototype) = nothing, nothing