Fixes

Author: gdalle

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/forward_onearg.jl

@@ -46,7 +46,7 @@ function DI.value_and_pushforward(
         dy = _copy_output(last(y_and_dy))
         return y, dy
     end
-    y = first(ys_and_ty[1])
+    y = _copy_output(first(ys_and_ty[1]))
     ty = map(last, ys_and_ty)
     return y, ty
 end

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/forward_twoarg.jl

@@ -19,6 +19,7 @@ function DI.prepare_pushforward_nokwarg(
     _sig = DI.signature(f!, y, backend, x, tx, contexts...; strict)
     config = get_config(backend)
     cache = prepare_derivative_cache(
+        call_and_return,
         f!,
         y,
         x,
@@ -43,14 +44,15 @@ function DI.value_and_pushforward(
     DI.check_prep(f!, y, prep, backend, x, tx, contexts...)
     ty = map(tx) do dx
         dy = zero_tangent(y)  # TODO: remove allocation?
-        value_and_derivative!!(
+        _, new_dy = value_and_derivative!!(
             prep.cache,
+            (call_and_return, zero_tangent(call_and_return)),
             (f!, prep.df!),
             (y, dy),
             (x, dx),
             map(first_unwrap, contexts, prep.context_tangents)...,
         )
-        return dy
+        return _copy_output(new_dy)
     end
     return y, ty
 end
@@ -80,13 +82,15 @@ function DI.value_and_pushforward!(
     ) where {F, C, X, Y}
     DI.check_prep(f!, y, prep, backend, x, tx, contexts...)
     foreach(tx, ty) do dx, dy
-        value_and_derivative!!(
+        _, new_dy = value_and_derivative!!(
             prep.cache,
+            (call_and_return, zero_tangent(call_and_return)),
             (f!, prep.df!),
             (y, dy),
             (x, dx),
             map(first_unwrap, contexts, prep.context_tangents)...,
         )
+        copyto!(dy, new_dy)
     end
     return y, ty
 end

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/twoarg.jl

@@ -1,8 +1,7 @@
-struct MooncakeTwoArgPullbackPrep{SIG, Tcache, DY, F, N} <: DI.PullbackPrep{SIG}
+struct MooncakeTwoArgPullbackPrep{SIG, Tcache, DY, N} <: DI.PullbackPrep{SIG}
     _sig::Val{SIG}
     cache::Tcache
     dy_backup::DY
-    target_function::F
     args_to_zero::NTuple{N, Bool}
 end
 
@@ -16,13 +15,9 @@ function DI.prepare_pullback_nokwarg(
         contexts::Vararg{DI.Context, C}
     ) where {F, C}
     _sig = DI.signature(f!, y, backend, x, ty, contexts...; strict)
-    target_function = function (f!, y, x, contexts...)
-        f!(y, x, contexts...)
-        return y
-    end
     config = get_config(backend)
     cache = prepare_pullback_cache(
-        target_function,
+        call_and_return,
         f!,
         y,
         x,
@@ -32,14 +27,14 @@ function DI.prepare_pullback_nokwarg(
     dy_backup_after = zero_tangent(y)
     contexts_tup_false = map(_ -> false, contexts)
     args_to_zero = (
-        false,  # target_function
+        false,  # call_and_return
         false,  # f!
         false,  # y
         true,  # x
         contexts_tup_false...,
     )
     prep = MooncakeTwoArgPullbackPrep(
-        _sig, cache, dy_backup_after, target_function, args_to_zero
+        _sig, cache, dy_backup_after, args_to_zero
     )
     return prep
 end
@@ -61,7 +56,7 @@ function DI.value_and_pullback(
     y_after, (_, _, _, dx) = value_and_pullback!!(
         prep.cache,
         dy_backup_after,
-        prep.target_function,
+        call_and_return,
         f!,
         y,
         x,
@@ -87,7 +82,7 @@ function DI.value_and_pullback(
         y_after, (_, _, _, dx) = value_and_pullback!!(
             prep.cache,
             dy_backup_after,
-            prep.target_function,
+            call_and_return,
             f!,
             y,
             x,

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/utils.jl

@@ -3,3 +3,8 @@ get_config(backend::AnyAutoMooncake{<:Config}) = backend.config
 
 @inline zero_tangent_unwrap(c::DI.Context) = zero_tangent(DI.unwrap(c))
 @inline first_unwrap(c, dc) = (DI.unwrap(c), dc)
+
+function call_and_return(f!::F, y, x, contexts...) where {F}
+    f!(y, x, contexts...)
+    return y
+end

DifferentiationInterface/test/Back/Mooncake/Project.toml

@@ -7,6 +7,3 @@ Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
 SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
-
-[sources]
-Mooncake = { url = "https://github.com/gdalle/Mooncake.jl", rev = "gd/forwardcache" }