Small fixes to svd_trunc

Author: kshyatt

ext/MatrixAlgebraKitMooncakeExt/MatrixAlgebraKitMooncakeExt.jl

@@ -3,7 +3,7 @@ module MatrixAlgebraKitMooncakeExt
 using Mooncake
 using Mooncake: DefaultCtx, CoDual, Dual, NoRData, rrule!!, frule!!, arrayify, @is_primitive
 using MatrixAlgebraKit
-using MatrixAlgebraKit: inv_safe, diagview, copy_input, zero!
+using MatrixAlgebraKit: inv_safe, diagview, copy_input, zero!, truncate, truncation_error!
 using MatrixAlgebraKit: qr_pullback!, qr_pushforward!, lq_pullback!, lq_pushforward!
 using MatrixAlgebraKit: qr_null_pullback!, qr_null_pushforward!, lq_null_pullback!, lq_null_pushforward!
 using MatrixAlgebraKit: eig_pullback!, eigh_pullback!, eig_trunc_pullback!, eigh_trunc_pullback!
@@ -419,14 +419,25 @@ function Mooncake.frule!!(::Dual{typeof(svd_trunc)}, A_dA::Dual, alg_dalg::Dual)
     # compute primal
     A, dA  = Mooncake.arrayify(A_dA)
     alg    = Mooncake.primal(alg_dalg)
-    output = svd_trunc(A, alg)
+    USVᴴ   = svd_compact(A, alg.alg)
+    U, S, Vᴴ = USVᴴ
+    dUfull  = zeros(eltype(U), size(U))
+    dSfull  = Diagonal(zeros(eltype(S), length(diagview(S))))
+    dVᴴfull = zeros(eltype(Vᴴ), size(Vᴴ))
+    svd_pushforward!(dA, A, (U, S, Vᴴ), (dUfull, dSfull, dVᴴfull))
+
+    USVᴴtrunc, ind = truncate(svd_trunc!, USVᴴ, alg.trunc)
+    ϵ        = truncation_error!(diagview(S), ind)
+    output   = (USVᴴtrunc..., ϵ)
     output_dual = Mooncake.zero_dual(output)
-    U, S, Vᴴ, ϵ = output
+    Utrunc, Strunc, Vᴴtrunc, ϵ = output
     dU_, dS_, dVᴴ_, dϵ = Mooncake.tangent(output_dual)
-    dU = arrayify(U, dU_)
-    dS = arrayify(S, dS_)
-    dVᴴ = arrayify(Vᴴ, dVᴴ_)
-    svd_trunc_pushforward!(dA, A, (U, S, Vᴴ), (dU, dS, dVᴴ))
+    Utrunc, dU  = arrayify(Utrunc, dU_)
+    Strunc, dS  = arrayify(Strunc, dS_)
+    Vᴴtrunc, dVᴴ = arrayify(Vᴴtrunc, dVᴴ_)
+    dU .= view(dUfull, :, ind)
+    diagview(dS) .= view(diagview(dSfull), ind)
+    dVᴴ .= view(dVᴴfull, ind, :)
     return output_dual
 end
 

src/pushforwards/eigh.jl

@@ -1,16 +1,18 @@
 function eigh_pushforward!(dA, A, DV, dDV; kwargs...)
-    D, V         = DV
-    dD, dV       = dDV
-    tmpV         = V \ dA
-    ∂K           = tmpV * V
-    ∂Kdiag       = diag(∂K)
-    dD.diag     .= real.(∂Kdiag)
-    dDD          = transpose(diagview(D)) .- diagview(D)
-    F            = one(eltype(dDD)) ./ dDD
-    diagview(F) .= zero(eltype(F))
-    ∂K         .*= F
-    ∂V           = mul!(tmpV, V, ∂K) 
-    copyto!(dV, ∂V)
+    D, V          = DV
+    dD, dV        = dDV
+    tmpV          = V \ dA
+    ∂K            = tmpV * V
+    ∂Kdiag        = diag(∂K)
+    diagview(dD) .= real.(∂Kdiag)
+    if !iszerotangent(dV)
+        dDD          = transpose(diagview(D)) .- diagview(D)
+        F            = one(eltype(dDD)) ./ dDD
+        diagview(F) .= zero(eltype(F))
+        ∂K         .*= F
+        ∂V           = mul!(tmpV, V, ∂K) 
+        copyto!(dV, ∂V)
+    end
     return (dD, dV)
 end
 

src/pushforwards/qr.jl

@@ -50,7 +50,7 @@ function qr_pushforward!(dA, A, QR, dQR; tol::Real=default_pullback_gauge_atol(Q
         dQ3  .= Q3
     end
     if !isempty(dR22)
-        _, r22 = qr_full(dA2 - dQ1*R12 - Q1*dR12, LAPACK_HouseholderQR(; positive=true))
+        _, r22 = qr_compact(dA2 - dQ1*R12 - Q1*dR12; positive=true)
         dR22  .= view(r22, 1:size(dR22, 1), 1:size(dR22, 2))
     end
     return (dQ, dR)

src/pushforwards/svd.jl

@@ -77,4 +77,6 @@ function svd_pushforward!(ΔA, A, USVᴴ, ΔUSVᴴ; rank_atol=default_pullback_r
     return (ΔU, ΔS, ΔVᴴ)
 end
 
-function svd_trunc_pushforward!(ΔA, A, USVᴴ, ΔUSVᴴ; rank_atol=default_pullback_rank_atol(A), kwargs...) end
+function svd_trunc_pushforward!(ΔA, A, USVᴴ, ΔUSVᴴ, ind; rank_atol=default_pullback_rank_atol(A), kwargs...)
+
+end

test/mooncake.jl

@@ -151,9 +151,9 @@ end
                 dQ   = make_mooncake_tangent(copy(ΔQ))
                 dR   = make_mooncake_tangent(copy(ΔR))
                 dQR  = Mooncake.build_tangent(typeof((ΔQ, ΔR)), dQ, dR)
-                Mooncake.TestUtils.test_rule(rng, qr_full, A, alg; output_tangent = dQR, atol=atol, rtol=rtol)
-                #Mooncake.TestUtils.test_rule(rng, ((A, alg)->qr_full(A, alg)[2]), A, alg; mode=Mooncake.ForwardMode, is_primitive=false, atol=atol, rtol=rtol)
-                #Mooncake.TestUtils.test_rule(rng, ((A, alg)->qr_full(A, alg)[1][1:m, 1:minmn]), A, alg; mode=Mooncake.ForwardMode, is_primitive=false, atol=atol, rtol=rtol)
+                #Mooncake.TestUtils.test_rule(rng, qr_full, A, alg; output_tangent = dQR, atol=atol, rtol=rtol)
+                Mooncake.TestUtils.test_rule(rng, ((A, alg)->qr_full(A, alg)[2]), A, alg; mode=Mooncake.ForwardMode, is_primitive=false, atol=atol, rtol=rtol)
+                Mooncake.TestUtils.test_rule(rng, ((A, alg)->qr_full(A, alg)[1][1:m, 1:minmn]), A, alg; mode=Mooncake.ForwardMode, is_primitive=false, atol=atol, rtol=rtol)
                 Mooncake.TestUtils.test_rule(rng, ((A, alg)->qr_full(A, alg)[1][1:m, minmn+1:m]), A, alg; mode=Mooncake.ForwardMode, is_primitive=false, atol=atol, rtol=rtol)
                 test_pullbacks_match(rng, qr_full!, qr_full, A, (Q, R), (ΔQ, ΔR), alg)
             end
@@ -173,6 +173,11 @@ end
                 dR   = make_mooncake_tangent(copy(ΔR))
                 dQR  = Mooncake.build_tangent(typeof((ΔQ, ΔR)), dQ, dR)
                 Mooncake.TestUtils.test_rule(rng, qr_compact, copy(Ard), alg; output_tangent = dQR, atol=atol, rtol=rtol)
+                Mooncake.TestUtils.test_rule(rng, ((A, alg)->qr_compact(A, alg)[2]), A, alg; mode=Mooncake.ForwardMode, is_primitive=false, atol=atol, rtol=rtol)
+                Mooncake.TestUtils.test_rule(rng, ((A, alg)->qr_compact(A, alg)[1][1:r, 1:r]), A, alg; mode=Mooncake.ForwardMode, is_primitive=false, atol=atol, rtol=rtol)
+                Mooncake.TestUtils.test_rule(rng, ((A, alg)->qr_compact(A, alg)[1][r+1:m, 1:r]), A, alg; mode=Mooncake.ForwardMode, is_primitive=false, atol=atol, rtol=rtol)
+                Mooncake.TestUtils.test_rule(rng, ((A, alg)->qr_compact(A, alg)[1][1:r, r+1:minmn]), A, alg; mode=Mooncake.ForwardMode, is_primitive=false, atol=atol, rtol=rtol)
+                Mooncake.TestUtils.test_rule(rng, ((A, alg)->qr_compact(A, alg)[1][r+1:m, r+1:minmn]), A, alg; mode=Mooncake.ForwardMode, is_primitive=false, atol=atol, rtol=rtol)
                 test_pullbacks_match(rng, qr_compact!, qr_compact, Ard, (Q, R), (ΔQ, ΔR), alg)
             end
         end