Fixes

Author: gdalle

DifferentiationInterface/ext/DifferentiationInterfaceEnzymeExt/forward_onearg.jl

@@ -224,7 +224,7 @@ function DI.jacobian(
     x,
     contexts::Vararg{DI.Constant,C},
 ) where {F,B,C}
-    DI.check_prep(f, prep, backend, contexts...)
+    DI.check_prep(f, prep, backend, x, contexts...)
     mode = forward_noprimal(backend)
     f_and_df = get_f_and_df(f, backend, mode)
     annotated_contexts = translate(backend, mode, Val(B), contexts...)
@@ -242,7 +242,7 @@ function DI.value_and_jacobian(
     x,
     contexts::Vararg{DI.Constant,C},
 ) where {F,B,C}
-    DI.check_prep(f, prep, backend, contexts...)
+    DI.check_prep(f, prep, backend, x, contexts...)
     mode = forward_withprimal(backend)
     f_and_df = get_f_and_df(f, backend, mode)
     annotated_contexts = translate(backend, mode, Val(B), contexts...)
@@ -261,7 +261,7 @@ function DI.jacobian!(
     x,
     contexts::Vararg{DI.Constant,C},
 ) where {F,C}
-    DI.check_prep(f, prep, backend, contexts...)
+    DI.check_prep(f, prep, backend, x, contexts...)
     return copyto!(jac, DI.jacobian(f, prep, backend, x, contexts...))
 end
 
@@ -273,7 +273,7 @@ function DI.value_and_jacobian!(
     x,
     contexts::Vararg{DI.Constant,C},
 ) where {F,C}
-    DI.check_prep(f, prep, backend, contexts...)
+    DI.check_prep(f, prep, backend, x, contexts...)
     y, new_jac = DI.value_and_jacobian(f, prep, backend, x, contexts...)
     return y, copyto!(jac, new_jac)
 end

DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/twoarg.jl

@@ -48,8 +48,13 @@ function compute_ydual_twoarg(
     tx::NTuple{B},
     contexts::Vararg{DI.Context,C},
 ) where {F,SIG,T,B,C}
-    (; xdual_tmp, ydual_tmp) = prep
-    make_dual!(T, xdual_tmp, x, tx)
+    (; ydual_tmp) = prep
+    if DI.ismutable_array(x)
+        make_dual!(T, prep.xdual_tmp, x, tx)
+        xdual_tmp = prep.xdual_tmp
+    else
+        xdual_tmp = make_dual(T, x, tx)
+    end
     contexts_dual = translate_prepared(contexts, prep.contexts_dual)
     f!(ydual_tmp, xdual_tmp, contexts_dual...)
     return ydual_tmp

DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl

@@ -25,7 +25,7 @@ function DI.prepare_pushforward(
     end
     if x isa Number
         xt = TPS{promote_type(typeof(first(tx)), typeof(x), Float64)}(; use=d)
-        return GTPSAOneArgPushforwardPrep(xt)
+        return GTPSAOneArgPushforwardPrep(_sig, xt)
     else
         xt = similar(x, TPS{promote_type(eltype(first(tx)), eltype(x), Float64)})
         for i in eachindex(xt)

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/onearg.jl

@@ -43,33 +43,21 @@ function DI.value_and_pullback(
     return new_y, (mycopy(new_dx),)
 end
 
-function DI.value_and_pullback!(
-    f::F,
-    tx::NTuple{1},
-    prep::MooncakeOneArgPullbackPrep{Y},
-    backend::AutoMooncake,
-    x,
-    ty::NTuple{1},
-    contexts::Vararg{DI.Context,C},
-) where {F,Y,C}
-    DI.check_prep(f, prep, backend, x, ty, contexts...)
-    y, (new_dx,) = DI.value_and_pullback(f, prep, backend, x, ty, contexts...)
-    copyto!(only(tx), new_dx)
-    return y, tx
-end
-
 function DI.value_and_pullback(
     f::F,
-    prep::MooncakeOneArgPullbackPrep,
+    prep::MooncakeOneArgPullbackPrep{Y},
     backend::AutoMooncake,
     x,
     ty::NTuple,
     contexts::Vararg{DI.Context,C},
-) where {F,C}
+) where {F,Y,C}
     DI.check_prep(f, prep, backend, x, ty, contexts...)
     ys_and_tx = map(ty) do dy
-        y, tx = DI.value_and_pullback(f, prep, backend, x, (dy,), contexts...)
-        y, only(tx)
+        dy_righttype = dy isa tangent_type(Y) ? dy : copyto!!(prep.dy_righttype, dy)
+        y, (_, new_dx) = value_and_pullback!!(
+            prep.cache, dy_righttype, f, x, map(DI.unwrap, contexts)...
+        )
+        y, mycopy(new_dx)
     end
     y = first(ys_and_tx[1])
     tx = last.(ys_and_tx)
@@ -86,11 +74,8 @@ function DI.value_and_pullback!(
     contexts::Vararg{DI.Context,C},
 ) where {F,C}
     DI.check_prep(f, prep, backend, x, ty, contexts...)
-    ys = map(tx, ty) do dx, dy
-        y, _ = DI.value_and_pullback!(f, (dx,), prep, backend, x, (dy,), contexts...)
-        y
-    end
-    y = ys[1]
+    y, new_tx = DI.value_and_pullback(f, prep, backend, x, ty, contexts...)
+    foreach(copyto!, tx, new_tx)
     return y, tx
 end
 

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/twoarg.jl

@@ -45,14 +45,18 @@ function DI.value_and_pullback(
     contexts::Vararg{DI.Context,C},
 ) where {F,C}
     DI.check_prep(f!, y, prep, backend, x, ty, contexts...)
-    # Prepare cotangent to add after the forward pass.
     dy = only(ty)
+    # Prepare cotangent to add after the forward pass.
     dy_righttype_after = copyto!(prep.dy_righttype, dy)
-
     # Run the reverse-pass and return the results.
-    contexts = map(DI.unwrap, contexts)
     y_after, (_, _, _, dx) = value_and_pullback!!(
-        prep.cache, dy_righttype_after, prep.target_function, f!, y, x, contexts...
+        prep.cache,
+        dy_righttype_after,
+        prep.target_function,
+        f!,
+        y,
+        x,
+        map(DI.unwrap, contexts)...,
     )
     copyto!(y, y_after)
     return y, (mycopy(dx),)
@@ -69,8 +73,18 @@ function DI.value_and_pullback(
 ) where {F,C}
     DI.check_prep(f!, y, prep, backend, x, ty, contexts...)
     tx = map(ty) do dy
-        _, tx = DI.value_and_pullback(f!, y, prep, backend, x, (dy,), contexts...)
-        only(tx)
+        dy_righttype_after = copyto!(prep.dy_righttype, dy)
+        y_after, (_, _, _, dx) = value_and_pullback!!(
+            prep.cache,
+            dy_righttype_after,
+            prep.target_function,
+            f!,
+            y,
+            x,
+            map(DI.unwrap, contexts)...,
+        )
+        copyto!(y, y_after)
+        mycopy(dx)
     end
     return y, tx
 end