fix: convert raw Mooncake.Tangent in pullback/gradient results

On Julia 1.11, Mooncake may return raw Tangent objects instead of
friendly arrays for StaticArrays even with friendly_tangents=true.
Add _maybe_to_primal dispatch as a safety net that converts leaked
Tangent objects to primal-shaped values, no-op otherwise.

Author: sunxd3

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/onearg.jl

@@ -35,7 +35,7 @@ function DI.value_and_pullback(
     new_y, (_, new_dx) = value_and_pullback!!(
         prep.cache, dy, f, x, map(DI.unwrap, contexts)...; prep.args_to_zero
     )
-    return new_y, (_copy_output(new_dx),)
+    return new_y, (_maybe_to_primal(new_dx, x),)
 end
 
 function DI.value_and_pullback(
@@ -51,7 +51,7 @@ function DI.value_and_pullback(
         y, (_, new_dx) = value_and_pullback!!(
             prep.cache, dy, f, x, map(DI.unwrap, contexts)...; prep.args_to_zero
         )
-        y, _copy_output(new_dx)
+        y, _maybe_to_primal(new_dx, x)
     end
     y = first(ys_and_tx[1])
     tx = map(last, ys_and_tx)
@@ -134,7 +134,7 @@ function DI.value_and_gradient(
         prep.cache, f, x, map(DI.unwrap, contexts)...;
         prep.args_to_zero
     )
-    return y, _copy_output(new_grad)
+    return y, _maybe_to_primal(new_grad, x)
 end
 
 function DI.value_and_gradient!(
@@ -150,7 +150,7 @@ function DI.value_and_gradient!(
         prep.cache, f, x, map(DI.unwrap, contexts)...;
         prep.args_to_zero
     )
-    copyto!(grad, new_grad)
+    copyto!(grad, _maybe_to_primal(new_grad, x))
     return y, grad
 end
 

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/twoarg.jl

@@ -64,7 +64,7 @@ function DI.value_and_pullback(
         prep.args_to_zero
     )
     copyto!(y, y_after)
-    return y, (_copy_output(dx),)
+    return y, (_maybe_to_primal(dx, x),)
 end
 
 function DI.value_and_pullback(
@@ -90,7 +90,7 @@ function DI.value_and_pullback(
             prep.args_to_zero
         )
         copyto!(y, y_after)
-        _copy_output(dx)
+        _maybe_to_primal(dx, x)
     end
     return y, tx
 end

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/utils.jl

@@ -21,6 +21,11 @@ function zero_tangent_or_primal(x, backend::AnyAutoMooncake)
     end
 end
 
+# Safety net: if Mooncake returns a raw Tangent (e.g. Julia 1.11 + StaticArrays),
+# convert it to a primal-shaped value. No-op for already-converted results.
+_maybe_to_primal(tx, x) = _copy_output(tx)
+_maybe_to_primal(tx::Mooncake.Tangent, x) = tangent_to_user_primal(tx, x)
+
 @inline maybe_getfield(mod, name::Symbol) =
     isdefined(mod, name) ? getfield(mod, name) : nothing
 

DifferentiationInterface/test/Back/Mooncake/test.jl

@@ -75,15 +75,12 @@ test_differentiation(
     @test grad.B == ps.A
 end
 
-# friendly_tangents + StaticArrays broken on Julia 1.11 (upstream Mooncake bug)
-@static if !(VERSION ≥ v"1.11-" && VERSION < v"1.12-")
-    test_differentiation(
-        backends[3:4],
-        nomatrix(static_scenarios());
-        logging = LOGGING,
-        excluded = SECOND_ORDER,
-    )
-end
+test_differentiation(
+    backends[3:4],
+    nomatrix(static_scenarios());
+    logging = LOGGING,
+    excluded = SECOND_ORDER,
+)
 
 @testset "Friendly tangents structured matrices" begin
     backend = AutoMooncake(; config = Mooncake.Config(; friendly_tangents = true))