Update Mooncake support to 0.5

Author: Technici4n

DifferentiationInterface/Project.toml

@@ -69,7 +69,7 @@ ForwardDiff = "0.10.36,1"
 GPUArraysCore = "0.2"
 GTPSA = "1.4.0"
 LinearAlgebra = "1"
-Mooncake = "0.4.175"
+Mooncake = "0.5.0"
 PolyesterForwardDiff = "0.1.2"
 ReverseDiff = "1.15.1"
 SparseArrays = "1"

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/forward_onearg.jl

@@ -1,5 +1,7 @@
 ## Pushforward
 
+# TODO: needs friendly tangents support
+
 struct MooncakeOneArgPushforwardPrep{SIG, Tcache, DX, FT, CT} <: DI.PushforwardPrep{SIG}
     _sig::Val{SIG}
     cache::Tcache
@@ -19,7 +21,7 @@ function DI.prepare_pushforward_nokwarg(
     _sig = DI.signature(f, backend, x, tx, contexts...; strict)
     config = get_config(backend)
     cache = prepare_derivative_cache(
-        f, x, map(DI.unwrap, contexts)...; config.debug_mode, config.silence_debug_messages
+        f, x, map(DI.unwrap, contexts)...; config
     )
     dx_righttype = zero_tangent(x)
     df = zero_tangent(f)

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/forward_twoarg.jl

@@ -1,5 +1,7 @@
 ## Pushforward
 
+# TODO: needs friendly tangents support
+
 struct MooncakeTwoArgPushforwardPrep{SIG, Tcache, DX, DY, FT, CT} <: DI.PushforwardPrep{SIG}
     _sig::Val{SIG}
     cache::Tcache
@@ -25,8 +27,7 @@ function DI.prepare_pushforward_nokwarg(
         y,
         x,
         map(DI.unwrap, contexts)...;
-        config.debug_mode,
-        config.silence_debug_messages,
+        config,
     )
     dx_righttype = zero_tangent(x)
     dy_righttype = zero_tangent(y)

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/onearg.jl

@@ -1,9 +1,8 @@
 ## Pullback
 
-struct MooncakeOneArgPullbackPrep{SIG, Tcache, DY, N} <: DI.PullbackPrep{SIG}
+struct MooncakeOneArgPullbackPrep{SIG, Tcache, N} <: DI.PullbackPrep{SIG}
     _sig::Val{SIG}
     cache::Tcache
-    dy_righttype::DY
     args_to_zero::NTuple{N, Bool}
 end
 
@@ -12,18 +11,14 @@ function DI.prepare_pullback_nokwarg(
     ) where {F, C}
     _sig = DI.signature(f, backend, x, ty, contexts...; strict)
     config = get_config(backend)
-    cache = prepare_pullback_cache(
-        f, x, map(DI.unwrap, contexts)...; config.debug_mode, config.silence_debug_messages
-    )
-    y = f(x, map(DI.unwrap, contexts)...)
-    dy_righttype = zero_tangent(y)
+    cache = prepare_pullback_cache(f, x, map(DI.unwrap, contexts)...; config)
     contexts_tup_false = map(_ -> false, contexts)
     args_to_zero = (
         false,  # f
         true,  # x
         contexts_tup_false...,
     )
-    prep = MooncakeOneArgPullbackPrep(_sig, cache, dy_righttype, args_to_zero)
+    prep = MooncakeOneArgPullbackPrep(_sig, cache, args_to_zero)
     return prep
 end
 
@@ -37,9 +32,8 @@ function DI.value_and_pullback(
     ) where {F, Y, C}
     DI.check_prep(f, prep, backend, x, ty, contexts...)
     dy = only(ty)
-    dy_righttype = dy isa tangent_type(Y) ? dy : _copy_to_output!!(prep.dy_righttype, dy)
     new_y, (_, new_dx) = value_and_pullback!!(
-        prep.cache, dy_righttype, f, x, map(DI.unwrap, contexts)...;
+        prep.cache, dy, f, x, map(DI.unwrap, contexts)...;
         prep.args_to_zero
     )
     return new_y, (_copy_output(new_dx),)
@@ -55,10 +49,8 @@ function DI.value_and_pullback(
     ) where {F, Y, C}
     DI.check_prep(f, prep, backend, x, ty, contexts...)
     ys_and_tx = map(ty) do dy
-        dy_righttype =
-            dy isa tangent_type(Y) ? dy : _copy_to_output!!(prep.dy_righttype, dy)
         y, (_, new_dx) = value_and_pullback!!(
-            prep.cache, dy_righttype, f, x, map(DI.unwrap, contexts)...;
+            prep.cache, dy, f, x, map(DI.unwrap, contexts)...;
             prep.args_to_zero
         )
         y, _copy_output(new_dx)
@@ -121,9 +113,7 @@ function DI.prepare_gradient_nokwarg(
     ) where {F, C}
     _sig = DI.signature(f, backend, x, contexts...; strict)
     config = get_config(backend)
-    cache = prepare_gradient_cache(
-        f, x, map(DI.unwrap, contexts)...; config.debug_mode, config.silence_debug_messages
-    )
+    cache = prepare_gradient_cache(f, x, map(DI.unwrap, contexts)...; config)
     contexts_tup_false = map(_ -> false, contexts)
     args_to_zero = (
         false,  # f

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/twoarg.jl

@@ -1,7 +1,6 @@
-struct MooncakeTwoArgPullbackPrep{SIG, Tcache, DY, F, N} <: DI.PullbackPrep{SIG}
+struct MooncakeTwoArgPullbackPrep{SIG, Tcache, F, N} <: DI.PullbackPrep{SIG}
     _sig::Val{SIG}
     cache::Tcache
-    dy_righttype::DY
     target_function::F
     args_to_zero::NTuple{N, Bool}
 end
@@ -27,10 +26,8 @@ function DI.prepare_pullback_nokwarg(
         y,
         x,
         map(DI.unwrap, contexts)...;
-        debug_mode = config.debug_mode,
-        silence_debug_messages = config.silence_debug_messages,
+        config,
     )
-    dy_righttype_after = zero_tangent(y)
     contexts_tup_false = map(_ -> false, contexts)
     args_to_zero = (
         false,  # target_function
@@ -39,9 +36,7 @@ function DI.prepare_pullback_nokwarg(
         true,  # x
         contexts_tup_false...,
     )
-    prep = MooncakeTwoArgPullbackPrep(
-        _sig, cache, dy_righttype_after, target_function, args_to_zero
-    )
+    prep = MooncakeTwoArgPullbackPrep(_sig, cache, target_function, args_to_zero)
     return prep
 end
 
@@ -56,12 +51,10 @@ function DI.value_and_pullback(
     ) where {F, C}
     DI.check_prep(f!, y, prep, backend, x, ty, contexts...)
     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.
     y_after, (_, _, _, dx) = value_and_pullback!!(
         prep.cache,
-        dy_righttype_after,
+        dy,
         prep.target_function,
         f!,
         y,