simplify implementations

Author: lkdvos

ext/MatrixAlgebraKitChainRulesCoreExt.jl

@@ -274,46 +274,22 @@ 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 ChainRulesCore.rrule(::typeof(project_hermitian), A, alg)
+    Aₕ = project_hermitian(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()
+        return NoTangent(), ΔA, 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 ChainRulesCore.rrule(::typeof(project_antihermitian), A, alg)
+    Aₐ = project_antihermitian(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()
+        return NoTangent(), ΔA, 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/MatrixAlgebraKitMooncakeExt/MatrixAlgebraKitMooncakeExt.jl

@@ -779,82 +779,44 @@ 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}
+        @is_primitive 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)
+            arg, darg = A_dA === arg_darg ? (A, dA) : arrayify(arg_darg)
             argc = copy(arg)
-            $f!(A, arg, Mooncake.primal(alg_dalg))
+            arg = $f!(A, arg, Mooncake.primal(alg_dalg))
+
             function $adj(::NoRData)
-                copy!(A, Ac)
                 dA .+= $f(darg)
+                dA === darg || zero!(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}
+
+        @is_primitive 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))
+            output_doutput = Mooncake.zero_fcodual(output)
+
+            doutput = last(arrayify(output_doutput))
             function $adj(::NoRData)
-                arg, darg = arrayify(output_codual)
-                dA .+= $f(darg)
-                zero!(darg)
-                return NoRData(), NoRData(), NoRData()
+                # TODO: need accumulating projection to avoid intermediate here
+                dA .+= $f(doutput)
+                return ntuple(Returns(NoRData(), 3))
             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