Add AD rules for projection methods

Add rrules/pullbacks for `project_hermitian!`, `project_antihermitian!`,
and `project_isometric!` directly in each AD backend extension
(ChainRulesCore, Enzyme, Mooncake). The hermitian/antihermitian
pullbacks are self-adjoint, while the isometric pullback delegates to
`left_polar_pullback!` with zero ΔP.

Includes test utilities and tests for all three backends.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: lkdvos

ext/MatrixAlgebraKitChainRulesCoreExt.jl

@@ -274,4 +274,46 @@ function ChainRulesCore.rrule(::typeof(right_polar!), A, PWᴴ, alg)
     return PWᴴ, right_polar_pullback
 end
 
+function ChainRulesCore.rrule(::typeof(project_hermitian!), A, Aₕ, alg)
+    Ac = copy_input(project_hermitian, A)
+    Aₕ = project_hermitian!(Ac, Aₕ, alg)
+    function project_hermitian_pullback(ΔAₕ)
+        ΔA = project_hermitian(unthunk(ΔAₕ))
+        return NoTangent(), ΔA, ZeroTangent(), NoTangent()
+    end
+    function project_hermitian_pullback(::ZeroTangent)
+        return NoTangent(), ZeroTangent(), ZeroTangent(), NoTangent()
+    end
+    return Aₕ, project_hermitian_pullback
+end
+
+function ChainRulesCore.rrule(::typeof(project_antihermitian!), A, Aₐ, alg)
+    Ac = copy_input(project_antihermitian, A)
+    Aₐ = project_antihermitian!(Ac, Aₐ, alg)
+    function project_antihermitian_pullback(ΔAₐ)
+        ΔA = project_antihermitian(unthunk(ΔAₐ))
+        return NoTangent(), ΔA, ZeroTangent(), NoTangent()
+    end
+    function project_antihermitian_pullback(::ZeroTangent)
+        return NoTangent(), ZeroTangent(), ZeroTangent(), NoTangent()
+    end
+    return Aₐ, project_antihermitian_pullback
+end
+
+function ChainRulesCore.rrule(::typeof(project_isometric!), A, W, alg)
+    Ac = copy_input(project_isometric, A)
+    # Compute the full polar decomposition to cache P for the pullback
+    WP = left_polar!(Ac, (similar(W), similar(W, size(W, 2), size(W, 2))), alg)
+    W_out = copy!(W, WP[1])
+    function project_isometric_pullback(ΔW)
+        ΔA = zero(A)
+        MatrixAlgebraKit.left_polar_pullback!(ΔA, A, WP, (unthunk(ΔW), nothing))
+        return NoTangent(), ΔA, ZeroTangent(), NoTangent()
+    end
+    function project_isometric_pullback(::ZeroTangent)
+        return NoTangent(), ZeroTangent(), ZeroTangent(), NoTangent()
+    end
+    return W_out, project_isometric_pullback
+end
+
 end

ext/MatrixAlgebraKitEnzymeExt/MatrixAlgebraKitEnzymeExt.jl

@@ -454,4 +454,91 @@ function EnzymeRules.reverse(
     return (nothing, nothing, nothing)
 end
 
+# single-output projections: project_hermitian!, project_antihermitian!
+# single-output projections: project_hermitian!, project_antihermitian!
+for (f!, project_f) in (
+        (project_hermitian!, project_hermitian),
+        (project_antihermitian!, project_antihermitian),
+    )
+    @eval begin
+        function EnzymeRules.augmented_primal(
+                config::EnzymeRules.RevConfigWidth{1},
+                func::Const{typeof($f!)},
+                ::Type{RT},
+                A::Annotation,
+                arg::Annotation{TA},
+                alg::Const{<:MatrixAlgebraKit.AbstractAlgorithm},
+            ) where {RT, TA}
+            ret = func.val(A.val, arg.val, alg.val)
+            cache_arg = (arg.val !== ret) || EnzymeRules.overwritten(config)[3] ? copy(ret) : nothing
+            dret = if EnzymeRules.needs_shadow(config)
+                (TA == Nothing || isa(arg, Const)) ? zero(ret) : arg.dval
+            else
+                nothing
+            end
+            primal = EnzymeRules.needs_primal(config) ? ret : nothing
+            return EnzymeRules.AugmentedReturn(primal, dret, (cache_arg, dret))
+        end
+        function EnzymeRules.reverse(
+                config::EnzymeRules.RevConfigWidth{1},
+                func::Const{typeof($f!)},
+                ::Type{RT},
+                cache,
+                A::Annotation,
+                arg::Annotation,
+                alg::Const{<:MatrixAlgebraKit.AbstractAlgorithm},
+            ) where {RT}
+            cache_arg, darg = cache
+            argdval = something(darg, arg.dval)
+            if !isa(A, Const)
+                A.dval .+= $project_f(argdval)
+            end
+            !isa(arg, Const) && make_zero!(arg.dval)
+            return (nothing, nothing, nothing)
+        end
+    end
+end
+
+# project_isometric! needs special handling: compute full polar decomposition
+function EnzymeRules.augmented_primal(
+        config::EnzymeRules.RevConfigWidth{1},
+        func::Const{typeof(project_isometric!)},
+        ::Type{RT},
+        A::Annotation,
+        W::Annotation{TW},
+        alg::Const{<:MatrixAlgebraKit.AbstractAlgorithm},
+    ) where {RT, TW}
+    # Compute the full polar decomposition for the pullback
+    Ac = copy(A.val)
+    m, n = size(A.val)
+    P = similar(A.val, n, n)
+    WP = left_polar!(Ac, (W.val, P), alg.val)
+    cache_WP = EnzymeRules.overwritten(config)[3] ? copy.(WP) : nothing
+    dret = if EnzymeRules.needs_shadow(config)
+        (TW == Nothing || isa(W, Const)) ? zero(WP[1]) : W.dval
+    else
+        nothing
+    end
+    primal = EnzymeRules.needs_primal(config) ? WP[1] : nothing
+    return EnzymeRules.AugmentedReturn(primal, dret, (cache_WP, dret))
+end
+function EnzymeRules.reverse(
+        config::EnzymeRules.RevConfigWidth{1},
+        func::Const{typeof(project_isometric!)},
+        ::Type{RT},
+        cache,
+        A::Annotation,
+        W::Annotation,
+        alg::Const{<:MatrixAlgebraKit.AbstractAlgorithm},
+    ) where {RT}
+    cache_WP, dW = cache
+    Aval = nothing
+    WPval = something(cache_WP, (W.val, cache_WP[2]))
+    if !isa(A, Const)
+        left_polar_pullback!(A.dval, Aval, WPval, (dW, nothing))
+    end
+    !isa(W, Const) && make_zero!(W.dval)
+    return (nothing, nothing, nothing)
+end
+
 end

ext/MatrixAlgebraKitMooncakeExt/MatrixAlgebraKitMooncakeExt.jl

@@ -779,4 +779,82 @@ function Mooncake.rrule!!(::CoDual{typeof(svd_trunc_no_error)}, A_dA::CoDual, al
     return USVᴴtrunc_dUSVᴴtrunc, svd_trunc_adjoint
 end
 
+# single-output projections: project_hermitian!, project_antihermitian!
+# single-output projections: project_hermitian!, project_antihermitian!
+for (f!, f, adj) in (
+        (:project_hermitian!, :project_hermitian, :project_hermitian_adjoint),
+        (:project_antihermitian!, :project_antihermitian, :project_antihermitian_adjoint),
+    )
+    @eval begin
+        @is_primitive Mooncake.DefaultCtx Mooncake.ReverseMode Tuple{typeof($f!), Any, Any, MatrixAlgebraKit.AbstractAlgorithm}
+        function Mooncake.rrule!!(f_df::CoDual{typeof($f!)}, A_dA::CoDual, arg_darg::CoDual, alg_dalg::CoDual{<:MatrixAlgebraKit.AbstractAlgorithm})
+            A, dA = arrayify(A_dA)
+            Ac = copy(A)
+            arg, darg = arrayify(arg_darg)
+            argc = copy(arg)
+            $f!(A, arg, Mooncake.primal(alg_dalg))
+            function $adj(::NoRData)
+                copy!(A, Ac)
+                dA .+= $f(darg)
+                copy!(arg, argc)
+                zero!(darg)
+                return NoRData(), NoRData(), NoRData(), NoRData()
+            end
+            return arg_darg, $adj
+        end
+        @is_primitive Mooncake.DefaultCtx Mooncake.ReverseMode Tuple{typeof($f), Any, MatrixAlgebraKit.AbstractAlgorithm}
+        function Mooncake.rrule!!(f_df::CoDual{typeof($f)}, A_dA::CoDual, alg_dalg::CoDual{<:MatrixAlgebraKit.AbstractAlgorithm})
+            A, dA = arrayify(A_dA)
+            output = $f(A, Mooncake.primal(alg_dalg))
+            output_codual = CoDual(output, Mooncake.zero_tangent(output))
+            function $adj(::NoRData)
+                arg, darg = arrayify(output_codual)
+                dA .+= $f(darg)
+                zero!(darg)
+                return NoRData(), NoRData(), NoRData()
+            end
+            return output_codual, $adj
+        end
+    end
+end
+
+# project_isometric! needs special handling: compute full polar decomposition
+@is_primitive Mooncake.DefaultCtx Mooncake.ReverseMode Tuple{typeof(project_isometric!), Any, Any, MatrixAlgebraKit.AbstractAlgorithm}
+function Mooncake.rrule!!(f_df::CoDual{typeof(project_isometric!)}, A_dA::CoDual, W_dW::CoDual, alg_dalg::CoDual{<:MatrixAlgebraKit.AbstractAlgorithm})
+    A, dA = arrayify(A_dA)
+    W, dW = arrayify(W_dW)
+    Ac = copy(A)
+    Wc = copy(W)
+    # Compute the full polar decomposition for the pullback
+    m, n = size(A)
+    P = similar(A, n, n)
+    WP = left_polar!(copy(A), (copy(W), P), Mooncake.primal(alg_dalg))
+    copy!(W, WP[1])
+    function project_isometric_adjoint(::NoRData)
+        copy!(A, Ac)
+        left_polar_pullback!(dA, A, WP, (dW, nothing))
+        copy!(W, Wc)
+        zero!(dW)
+        return NoRData(), NoRData(), NoRData(), NoRData()
+    end
+    return W_dW, project_isometric_adjoint
+end
+
+@is_primitive Mooncake.DefaultCtx Mooncake.ReverseMode Tuple{typeof(project_isometric), Any, MatrixAlgebraKit.AbstractAlgorithm}
+function Mooncake.rrule!!(f_df::CoDual{typeof(project_isometric)}, A_dA::CoDual, alg_dalg::CoDual{<:MatrixAlgebraKit.AbstractAlgorithm})
+    A, dA = arrayify(A_dA)
+    alg = Mooncake.primal(alg_dalg)
+    # Compute the full polar decomposition for the pullback
+    WP = left_polar(A, alg)
+    W_out = WP[1]
+    output_codual = CoDual(W_out, Mooncake.zero_tangent(W_out))
+    function project_isometric_adjoint(::NoRData)
+        W, dW = arrayify(output_codual)
+        left_polar_pullback!(dA, A, WP, (dW, nothing))
+        zero!(dW)
+        return NoRData(), NoRData(), NoRData()
+    end
+    return output_codual, project_isometric_adjoint
+end
+
 end

test/testsuite/ad_utils.jl

@@ -421,3 +421,27 @@ function ad_right_null_setup(A)
     ΔNᴴ = randn!(similar(A, T, n - min(m, n), min(m, n))) * right_orth(A; alg = :lq)[2]
     return Nᴴ, ΔNᴴ
 end
+
+function ad_project_hermitian_setup(A)
+    m, n = size(A)
+    T = eltype(A)
+    Aₕ = project_hermitian(A)
+    ΔAₕ = randn!(similar(A, T, m, n))
+    return Aₕ, ΔAₕ
+end
+
+function ad_project_antihermitian_setup(A)
+    m, n = size(A)
+    T = eltype(A)
+    Aₐ = project_antihermitian(A)
+    ΔAₐ = randn!(similar(A, T, m, n))
+    return Aₐ, ΔAₐ
+end
+
+function ad_project_isometric_setup(A)
+    m, n = size(A)
+    T = eltype(A)
+    W = project_isometric(A)
+    ΔW = randn!(similar(A, T, m, n))
+    return W, ΔW
+end

test/testsuite/chainrules.jl

@@ -10,6 +10,7 @@ for f in
         :eig_trunc_no_error, :eigh_trunc_no_error,
         :svd_compact, :svd_trunc, :svd_trunc_no_error, :svd_vals,
         :left_polar, :right_polar,
+        :project_hermitian, :project_antihermitian, :project_isometric,
     )
     copy_f = Symbol(:cr_copy_, f)
     f! = Symbol(f, '!')
@@ -46,6 +47,7 @@ function test_chainrules(T::Type, sz; kwargs...)
         test_chainrules_svd(T, sz; kwargs...)
         test_chainrules_polar(T, sz; kwargs...)
         test_chainrules_orthnull(T, sz; kwargs...)
+        test_chainrules_projections(T, sz; kwargs...)
     end
 end
 
@@ -610,3 +612,58 @@ function test_chainrules_orthnull(
         )
     end
 end
+
+function test_chainrules_projections(
+        T::Type, sz;
+        atol::Real = 0, rtol::Real = precision(T),
+        kwargs...
+    )
+    summary_str = testargs_summary(T, sz)
+    return @testset "Projections Chainrules AD rules $summary_str" begin
+        A = instantiate_matrix(T, sz)
+        m, n = size(A)
+        config = Zygote.ZygoteRuleConfig()
+        if m == n
+            alg_h = MatrixAlgebraKit.default_hermitian_algorithm(A)
+            @testset "project_hermitian" begin
+                Aₕ, ΔAₕ = ad_project_hermitian_setup(A)
+                test_rrule(
+                    cr_copy_project_hermitian, A, alg_h ⊢ NoTangent();
+                    output_tangent = ΔAₕ, atol = atol, rtol = rtol
+                )
+                test_rrule(
+                    config, project_hermitian, A;
+                    output_tangent = ΔAₕ,
+                    atol = atol, rtol = rtol, rrule_f = rrule_via_ad, check_inferred = false
+                )
+            end
+            @testset "project_antihermitian" begin
+                Aₐ, ΔAₐ = ad_project_antihermitian_setup(A)
+                test_rrule(
+                    cr_copy_project_antihermitian, A, alg_h ⊢ NoTangent();
+                    output_tangent = ΔAₐ, atol = atol, rtol = rtol
+                )
+                test_rrule(
+                    config, project_antihermitian, A;
+                    output_tangent = ΔAₐ,
+                    atol = atol, rtol = rtol, rrule_f = rrule_via_ad, check_inferred = false
+                )
+            end
+        end
+        if m > n
+            @testset "project_isometric" begin
+                W, ΔW = ad_project_isometric_setup(A)
+                alg_iso = MatrixAlgebraKit.default_polar_algorithm(A)
+                test_rrule(
+                    cr_copy_project_isometric, A, alg_iso ⊢ NoTangent();
+                    output_tangent = ΔW, atol = atol, rtol = rtol
+                )
+                test_rrule(
+                    config, project_isometric, A;
+                    output_tangent = ΔW,
+                    atol = atol, rtol = rtol, rrule_f = rrule_via_ad, check_inferred = false
+                )
+            end
+        end
+    end
+end

test/testsuite/enzyme.jl

@@ -105,6 +105,7 @@ function test_enzyme(T::Type, sz; kwargs...)
             test_enzyme_polar(T, sz; kwargs...)
             test_enzyme_orthnull(T, sz; kwargs...)
         end
+        test_enzyme_projections(T, sz; kwargs...)
     end
 end
 
@@ -462,3 +463,41 @@ function test_enzyme_orthnull(
         end
     end
 end
+
+function test_enzyme_projections(
+        T::Type, sz;
+        atol::Real = 0, rtol::Real = precision(T),
+        kwargs...
+    )
+    summary_str = testargs_summary(T, sz)
+    return @testset "Projections Enzyme AD rules $summary_str" begin
+        A = instantiate_matrix(T, sz)
+        m, n = size(A)
+        fdm = eltype(T) <: Union{Float32, ComplexF32} ? EnzymeTestUtils.FiniteDifferences.central_fdm(5, 1, max_range = 1.0e-2) : EnzymeTestUtils.FiniteDifferences.central_fdm(5, 1)
+        if m == n
+            @testset "project_hermitian" begin
+                @testset "reverse: RT $RT, TA $TA" for RT in (Duplicated,), TA in (Duplicated,)
+                    Aₕ, ΔAₕ = ad_project_hermitian_setup(A)
+                    eltype(T) <: BlasFloat && test_reverse(project_hermitian, RT, (A, TA); atol, rtol, output_tangent = ΔAₕ, fdm)
+                    is_cpu(A) && enz_test_pullbacks_match(rng, project_hermitian!, project_hermitian, A, Aₕ, ΔAₕ)
+                end
+            end
+            @testset "project_antihermitian" begin
+                @testset "reverse: RT $RT, TA $TA" for RT in (Duplicated,), TA in (Duplicated,)
+                    Aₐ, ΔAₐ = ad_project_antihermitian_setup(A)
+                    eltype(T) <: BlasFloat && test_reverse(project_antihermitian, RT, (A, TA); atol, rtol, output_tangent = ΔAₐ, fdm)
+                    is_cpu(A) && enz_test_pullbacks_match(rng, project_antihermitian!, project_antihermitian, A, Aₐ, ΔAₐ)
+                end
+            end
+        end
+        if m > n
+            @testset "project_isometric" begin
+                @testset "reverse: RT $RT, TA $TA" for RT in (Duplicated,), TA in (Duplicated,)
+                    W, ΔW = ad_project_isometric_setup(A)
+                    eltype(T) <: BlasFloat && test_reverse(project_isometric, RT, (A, TA); atol, rtol, output_tangent = ΔW, fdm)
+                    is_cpu(A) && enz_test_pullbacks_match(rng, project_isometric!, project_isometric, A, W, ΔW)
+                end
+            end
+        end
+    end
+end