Working frule for polar

Author: kshyatt

ext/MatrixAlgebraKitMooncakeExt/MatrixAlgebraKitMooncakeExt.jl

@@ -216,7 +216,7 @@ for (f, St) in ((svd_full!, :AbstractMatrix), (svd_compact!, :Diagonal))
                     vS    = Diagonal(diagview(S)[1:minmn])
                     vVᴴ   = view(Vᴴ, 1:minmn, :)
                     vdU   = view(dU, :, 1:minmn)
-                    vdS   = Diagonal(diagview(dS)[1:minmn])
+                    vdS   = view(dS, 1:minmn, 1:minmn)
                     vdVᴴ  = view(dVᴴ, 1:minmn, :)
                     dA    = MatrixAlgebraKit.svd_pullback!(dA, A, (vU, vS, vVᴴ), (vdU, vdS, vdVᴴ))
                 end
@@ -243,7 +243,18 @@ for (f, St) in ((svd_full!, :AbstractMatrix), (svd_compact!, :Diagonal))
             U, dU          = arrayify(U_, dU_) 
             S, dS          = arrayify(S_, dS_) 
             Vᴴ, dVᴴ        = arrayify(Vᴴ_, dVᴴ_) 
-            (dU, dS, dVᴴ)  = svd_pushforward!(dA, A, (U, S, Vᴴ), (dU, dS, dVᴴ); kwargs...)
+            minmn          = min(size(A)...)
+            if ($f == svd_compact!) # compact
+                svd_pushforward!(dA, A, (U, S, Vᴴ), (dU, dS, dVᴴ); kwargs...)
+            else # full
+                vU    = view(U, :, 1:minmn)
+                vS    = view(S, 1:minmn, 1:minmn)
+                vVᴴ   = view(Vᴴ, 1:minmn, :)
+                vdU   = view(dU, :, 1:minmn)
+                vdS   = view(dS, 1:minmn, 1:minmn)
+                vdVᴴ  = view(dVᴴ, 1:minmn, :)
+                svd_pushforward!(dA, A, (vU, vS, vVᴴ), (vdU, vdS, vdVᴴ); kwargs...)
+            end
             return USVᴴ_dUSVᴴ
         end
     end

src/pushforwards/polar.jl

@@ -1,2 +1,25 @@
-function left_polar_pushforward! end
-function right_polar_pushforward! end
+function left_polar_pushforward!(ΔA, A, WP, ΔWP; kwargs...)
+    W, P   = WP
+    ΔW, ΔP = ΔWP
+    aWdA   = adjoint(W) * ΔA
+    K̇      = sylvester(P, P, -(aWdA - adjoint(aWdA)))
+    L̇      = (Diagonal(ones(eltype(W), size(W, 1))) - W*adjoint(W))*ΔA*inv(P)
+    ΔW     .= W * K̇ + L̇
+    ΔP     .= aWdA - K̇*P
+    MatrixAlgebraKit.zero!(ΔA)
+    return (ΔW, ΔP)
+end
+
+function right_polar_pushforward!(ΔA, A, PWᴴ, ΔPWᴴ; kwargs...)
+    P, Wᴴ   = PWᴴ
+    ΔP, ΔWᴴ = ΔPWᴴ
+    dAW     = ΔA * adjoint(Wᴴ)
+    K̇       = sylvester(P, P, -(dAW - adjoint(dAW)))
+    ImW     = (Diagonal(ones(eltype(Wᴴ), size(Wᴴ, 2))) - adjoint(Wᴴ) * Wᴴ)
+    @show size(P), size(ΔA), size(ImW), size(Wᴴ)
+    L̇       = inv(P)*ΔA*ImW
+    ΔWᴴ    .= K̇ * Wᴴ + L̇
+    ΔP     .= dAW - P * K̇
+    MatrixAlgebraKit.zero!(ΔA)
+    return (ΔWᴴ, ΔP)
+end