Fix and test friendly tangents with static arrays

Author: gdalle

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/DifferentiationInterfaceMooncakeExt.jl

@@ -29,7 +29,8 @@ using Mooncake:
     NoRData,
     primal,
     _copy_output,
-    _copy_to_output!!
+    _copy_to_output!!,
+    tangent_to_primal!!
 
 const AnyAutoMooncake{C} = Union{AutoMooncake{C}, AutoMooncakeForward{C}}
 

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/forward_onearg.jl

@@ -17,10 +17,8 @@ function DI.prepare_pushforward_nokwarg(
     ) where {F, C}
     _sig = DI.signature(f, backend, x, tx, contexts...; strict)
     config = get_config(backend)
-    cache = prepare_derivative_cache(
-        f, x, map(DI.unwrap, contexts)...; config
-    )
-    df = zero_tangent(f)
+    cache = prepare_derivative_cache(f, x, map(DI.unwrap, contexts)...; config)
+    df = zero_tangent_or_primal(f, backend)
     context_tangents = map(zero_tangent_unwrap, contexts)
     prep = MooncakeOneArgPushforwardPrep(_sig, cache, df, context_tangents)
     return prep

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/forward_twoarg.jl

@@ -26,7 +26,7 @@ function DI.prepare_pushforward_nokwarg(
         map(DI.unwrap, contexts)...;
         config
     )
-    df! = zero_tangent(f!)
+    df! = zero_tangent_or_primal(f!, backend)
     context_tangents = map(zero_tangent_unwrap, contexts)
     prep = MooncakeTwoArgPushforwardPrep(_sig, cache, df!, context_tangents)
     return prep
@@ -43,10 +43,11 @@ function DI.value_and_pushforward(
     ) where {F, C, X}
     DI.check_prep(f!, y, prep, backend, x, tx, contexts...)
     ty = map(tx) do dx
-        dy = zero_tangent(y)  # TODO: remove allocation?
+        dy = zero_tangent_or_primal(y, backend)  # TODO: remove allocation?
+        dcall = zero_tangent_or_primal(call_and_return, backend)
         _, new_dy = value_and_derivative!!(
             prep.cache,
-            (call_and_return, zero_tangent(call_and_return)),
+            (call_and_return, dcall),
             (f!, prep.df!),
             (y, dy),
             (x, dx),
@@ -82,9 +83,10 @@ 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
+        dcall = zero_tangent_or_primal(call_and_return, backend)
         _, new_dy = value_and_derivative!!(
             prep.cache,
-            (call_and_return, zero_tangent(call_and_return)),
+            (call_and_return, dcall),
             (f!, prep.df!),
             (y, dy),
             (x, dx),

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/onearg.jl

@@ -11,9 +11,7 @@ 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
-    )
+    cache = prepare_pullback_cache(f, x, map(DI.unwrap, contexts)...; config)
     contexts_tup_false = map(_ -> false, contexts)
     args_to_zero = (
         false,  # f
@@ -113,9 +111,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
-    )
+    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

@@ -24,7 +24,7 @@ function DI.prepare_pullback_nokwarg(
         map(DI.unwrap, contexts)...;
         config,
     )
-    dy_backup_after = zero_tangent(y)
+    dy_backup = zero_tangent_or_primal(y, backend)
     contexts_tup_false = map(_ -> false, contexts)
     args_to_zero = (
         false,  # call_and_return
@@ -34,7 +34,7 @@ function DI.prepare_pullback_nokwarg(
         contexts_tup_false...,
     )
     prep = MooncakeTwoArgPullbackPrep(
-        _sig, cache, dy_backup_after, args_to_zero
+        _sig, cache, dy_backup, args_to_zero
     )
     return prep
 end
@@ -51,11 +51,11 @@ function DI.value_and_pullback(
     DI.check_prep(f!, y, prep, backend, x, ty, contexts...)
     dy = only(ty)
     # Prepare cotangent to add after the forward pass.
-    dy_backup_after = copyto!(prep.dy_backup, dy)
+    dy_backup = copyto!(prep.dy_backup, dy)
     # Run the reverse-pass and return the results.
     y_after, (_, _, _, dx) = value_and_pullback!!(
         prep.cache,
-        dy_backup_after,
+        dy_backup,
         call_and_return,
         f!,
         y,
@@ -78,10 +78,10 @@ function DI.value_and_pullback(
     ) where {F, C}
     DI.check_prep(f!, y, prep, backend, x, ty, contexts...)
     tx = map(ty) do dy
-        dy_backup_after = copyto!(prep.dy_backup, dy)
+        dy_backup = copyto!(prep.dy_backup, dy)
         y_after, (_, _, _, dx) = value_and_pullback!!(
             prep.cache,
-            dy_backup_after,
+            dy_backup,
             call_and_return,
             f!,
             y,

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/utils.jl

@@ -8,3 +8,12 @@ function call_and_return(f!::F, y, x, contexts...) where {F}
     f!(y, x, contexts...)
     return y
 end
+
+function zero_tangent_or_primal(x, backend::AnyAutoMooncake)
+    if backend.config.friendly_tangents
+        # zero(x) but safer
+        return tangent_to_primal!!(_copy_output(x), zero_tangent(x))
+    else
+        return zero_tangent(x)
+    end
+end

DifferentiationInterface/test/Back/Mooncake/Project.toml

@@ -6,4 +6,5 @@ ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
 SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

DifferentiationInterface/test/Back/Mooncake/test.jl

@@ -9,8 +9,8 @@ check_no_implicit_imports(DifferentiationInterface)
 
 backends = [
     AutoMooncake(),
-    AutoMooncake(; config = Mooncake.Config(; friendly_tangents = true)),
     AutoMooncakeForward(),
+    AutoMooncake(; config = Mooncake.Config(; friendly_tangents = true)),
     AutoMooncakeForward(; config = Mooncake.Config(; friendly_tangents = true)),
 ]
 
@@ -22,7 +22,8 @@ end
 test_differentiation(
     backends,
     default_scenarios(;
-        include_constantified = true, include_cachified = true, use_tuples = true
+        include_batchified = false,
+        include_constantified = false, include_cachified = false, use_tuples = true
     );
     excluded = SECOND_ORDER,
     logging = LOGGING,
@@ -39,7 +40,7 @@ end
 
 # Test second-order differentiation (forward-over-reverse)
 test_differentiation(
-    [SecondOrder(AutoMooncakeForward(; config = nothing), AutoMooncake(; config = nothing))],
+    [SecondOrder(AutoMooncakeForward(), AutoMooncake())],
     excluded = EXCLUDED,
     logging = true,
 )
@@ -52,4 +53,9 @@ test_differentiation(
     @test grad.B == ps.A
 end
 
-# TODO: test static arrays with friendly tangents!
+test_differentiation(
+    backends[3:4],
+    static_scenarios();
+    logging = LOGGING,
+    excluded = SECOND_ORDER
+)