Enzyme

Author: gdalle

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/utils.jl

@@ -53,7 +53,9 @@ force_annotation(f::F) where {F} = Const(f)
     return Const(DI.unwrap(c))
 end
 
-@inline function _translate(backend::AutoEnzyme, ::Mode, ::Val{B}, c::DI.Cache) where {B}
+@inline function _translate(
+    backend::AutoEnzyme, ::Mode, ::Val{B}, c::Union{DI.Cache,DI.PrepContext}
+) where {B}
     if B == 1
         return Duplicated(DI.unwrap(c), make_zero(DI.unwrap(c)))
     else