fix Mooncake friendly_tangents compatibility

Author: sunxd3

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/utils.jl

@@ -11,9 +11,32 @@ end
 
 function zero_tangent_or_primal(x, backend::AnyAutoMooncake)
     if get_config(backend).friendly_tangents
-        # zero(x) but safer
-        return tangent_to_primal!!(_copy_output(x), zero_tangent(x))
+        # Mooncake 0.5.25+ replaced `tangent_to_primal!!` with the
+        # `tangent_to_friendly!!` framework. For this internal backup we still
+        # need a primal-shaped value, so use the `AsPrimal` path when
+        # available and fall back for older Mooncake releases.
+        return tangent_to_user_primal(zero_tangent(x), x)
     else
         return zero_tangent(x)
     end
 end
+
+@inline maybe_getfield(mod, name::Symbol) = isdefined(mod, name) ? getfield(mod, name) : nothing
+
+const mooncake_tangent_to_friendly = maybe_getfield(Mooncake, Symbol("tangent_to_friendly!!"))
+const mooncake_friendly_tangent_cache = maybe_getfield(Mooncake, :FriendlyTangentCache)
+const mooncake_as_primal = maybe_getfield(Mooncake, :AsPrimal)
+const mooncake_no_cache = maybe_getfield(Mooncake, :NoCache)
+
+function tangent_to_user_primal(tx, x)
+    if !isnothing(mooncake_tangent_to_friendly) &&
+       !isnothing(mooncake_friendly_tangent_cache) &&
+       !isnothing(mooncake_as_primal) &&
+       !isnothing(mooncake_no_cache)
+        dest = mooncake_friendly_tangent_cache{mooncake_as_primal}(_copy_output(x))
+        cache = isbitstype(typeof(x)) ? mooncake_no_cache() : IdDict{Any,Any}()
+        return mooncake_tangent_to_friendly(dest, x, tx, cache)
+    else
+        return tangent_to_primal!!(_copy_output(x), tx)
+    end
+end

DifferentiationInterface/test/Back/Mooncake/test.jl

@@ -1,6 +1,7 @@
 include("../../testutils.jl")
 
 using DifferentiationInterface, DifferentiationInterfaceTest
+using LinearAlgebra: Hermitian, SymTridiagonal, Symmetric
 using Mooncake: Mooncake
 using Test
 
@@ -80,3 +81,34 @@ test_differentiation(
     logging = LOGGING,
     excluded = SECOND_ORDER
 )
+
+@testset "Friendly tangents structured matrices" begin
+    backend = AutoMooncake(; config = Mooncake.Config(; friendly_tangents = true))
+    inputs = (
+        Symmetric([2.0 1.0; 1.0 3.0]),
+        Hermitian(ComplexF64[2 1 + im; 1 - im 3]),
+        SymTridiagonal([2.0, 3.0, 4.0], [5.0, 6.0]),
+    )
+    f(x) = real(sum(abs2, x))
+
+    @testset "$(typeof(x))" for x in inputs
+        grad = gradient(f, backend, x)
+        y, grad2 = value_and_gradient(f, backend, x)
+        pb = only(pullback(identity, backend, x, (x,)))
+
+        @test grad isa Matrix
+        @test grad2 isa Matrix
+        @test pb isa Matrix
+        @test grad == grad2
+        @test y == f(x)
+        @test pb == Matrix(x)
+
+        grad_dense = zero(Matrix(x))
+        @test gradient!(f, grad_dense, backend, x) === grad_dense
+        @test grad_dense == grad
+
+        tx_dense = (zero(Matrix(x)),)
+        @test only(pullback!(identity, tx_dense, backend, x, (x,))) === tx_dense[1]
+        @test tx_dense[1] == pb
+    end
+end