Fixes

Author: gdalle

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/forward_onearg.jl

@@ -137,11 +137,11 @@ end
 
 function DI.gradient(
     f::F,
-    prep::EnzymeForwardGradientPrep{B},
+    prep::EnzymeForwardGradientPrep{SIG,B},
     backend::AutoEnzyme{<:ForwardMode,<:Union{Nothing,Const}},
     x,
     contexts::Vararg{DI.Constant,C},
-) where {F,B,C}
+) where {F,SIG,B,C}
     DI.check_prep(f, prep, backend, x, contexts...)
     mode = forward_noprimal(backend)
     f_and_df = get_f_and_df(f, backend, mode)
@@ -154,11 +154,11 @@ end
 
 function DI.value_and_gradient(
     f::F,
-    prep::EnzymeForwardGradientPrep{B},
+    prep::EnzymeForwardGradientPrep{SIG,B},
     backend::AutoEnzyme{<:ForwardMode,<:Union{Nothing,Const}},
     x,
     contexts::Vararg{DI.Constant,C},
-) where {F,B,C}
+) where {F,SIG,B,C}
     DI.check_prep(f, prep, backend, x, contexts...)
     mode = forward_withprimal(backend)
     f_and_df = get_f_and_df(f, backend, mode)
@@ -172,23 +172,23 @@ end
 function DI.gradient!(
     f::F,
     grad,
-    prep::EnzymeForwardGradientPrep{B},
+    prep::EnzymeForwardGradientPrep{SIG,B},
     backend::AutoEnzyme{<:ForwardMode,<:Union{Nothing,Const}},
     x,
     contexts::Vararg{DI.Constant,C},
-) where {F,B,C}
+) where {F,SIG,B,C}
     DI.check_prep(f, prep, backend, x, contexts...)
     return copyto!(grad, DI.gradient(f, prep, backend, x, contexts...))
 end
 
 function DI.value_and_gradient!(
     f::F,
     grad,
-    prep::EnzymeForwardGradientPrep{B},
+    prep::EnzymeForwardGradientPrep{SIG,B},
     backend::AutoEnzyme{<:ForwardMode,<:Union{Nothing,Const}},
     x,
     contexts::Vararg{DI.Constant,C},
-) where {F,B,C}
+) where {F,SIG,B,C}
     DI.check_prep(f, prep, backend, x, contexts...)
     y, new_grad = DI.value_and_gradient(f, prep, backend, x, contexts...)
     return y, copyto!(grad, new_grad)
@@ -219,11 +219,11 @@ end
 
 function DI.jacobian(
     f::F,
-    prep::EnzymeForwardOneArgJacobianPrep{B},
+    prep::EnzymeForwardOneArgJacobianPrep{SIG,B},
     backend::AutoEnzyme{<:Union{ForwardMode,Nothing},<:Union{Nothing,Const}},
     x,
     contexts::Vararg{DI.Constant,C},
-) where {F,B,C}
+) where {F,SIG,B,C}
     DI.check_prep(f, prep, backend, x, contexts...)
     mode = forward_noprimal(backend)
     f_and_df = get_f_and_df(f, backend, mode)
@@ -237,11 +237,11 @@ end
 
 function DI.value_and_jacobian(
     f::F,
-    prep::EnzymeForwardOneArgJacobianPrep{B},
+    prep::EnzymeForwardOneArgJacobianPrep{SIG,B},
     backend::AutoEnzyme{<:Union{ForwardMode,Nothing},<:Union{Nothing,Const}},
     x,
     contexts::Vararg{DI.Constant,C},
-) where {F,B,C}
+) where {F,SIG,B,C}
     DI.check_prep(f, prep, backend, x, contexts...)
     mode = forward_withprimal(backend)
     f_and_df = get_f_and_df(f, backend, mode)

DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl

@@ -87,9 +87,8 @@ function DI.value_and_pushforward(
     contexts::Vararg{DI.Constant,C},
 ) where {C}
     DI.check_prep(f, prep, backend, x, tx, contexts...)
-    fc = DI.with_contexts(f, contexts...)
-    ty = DI.pushforward(fc, prep, backend, x, tx)
-    y = fc(x) # TO-DO: optimize
+    ty = DI.pushforward(f, prep, backend, x, tx, contexts...)
+    y = f(x, map(DI.unwrap, contexts)...)  # TODO: optimize
     return y, ty
 end
 
@@ -103,9 +102,8 @@ function DI.value_and_pushforward!(
     contexts::Vararg{DI.Constant,C},
 ) where {C}
     DI.check_prep(f, prep, backend, x, tx, contexts...)
-    fc = DI.with_contexts(f, contexts...)
-    DI.pushforward!(fc, ty, prep, backend, x, tx)
-    y = fc(x)  # TO-DO: optimize
+    DI.pushforward!(f, ty, prep, backend, x, tx, contexts...)
+    y = f(x, map(DI.unwrap, contexts)...)  # TODO: optimize
     return y, ty
 end
 

DifferentiationInterface/ext/DifferentiationInterfaceSymbolicsExt/onearg.jl

@@ -460,7 +460,7 @@ function DI.hvp!(
     tx::NTuple,
     contexts::Vararg{DI.Context,C},
 ) where {C}
-    DI.check_prep(f, prep, backend, x, contexts...)
+    DI.check_prep(f, prep, backend, x, tx, contexts...)
     for b in eachindex(tx, tg)
         dx, dg = tx[b], tg[b]
         prep.hvp_exe!(vec(dg), vec(x), vec(dx), map(DI.unwrap, contexts)...)
@@ -476,7 +476,7 @@ function DI.gradient_and_hvp(
     tx::NTuple,
     contexts::Vararg{DI.Context,C},
 ) where {C}
-    DI.check_prep(f, prep, backend, x, contexts...)
+    DI.check_prep(f, prep, backend, x, tx, contexts...)
     tg = DI.hvp(f, prep, backend, x, tx, contexts...)
     grad = DI.gradient(f, prep.gradient_prep, backend, x, contexts...)
     return grad, tg
@@ -492,7 +492,7 @@ function DI.gradient_and_hvp!(
     tx::NTuple,
     contexts::Vararg{DI.Context,C},
 ) where {C}
-    DI.check_prep(f, prep, backend, x, contexts...)
+    DI.check_prep(f, prep, backend, x, tx, contexts...)
     DI.hvp!(f, tg, prep, backend, x, tx, contexts...)
     DI.gradient!(f, grad, prep.gradient_prep, backend, x, contexts...)
     return grad, tg