Some fixes

Author: kshyatt

ext/MatrixAlgebraKitMooncakeExt/MatrixAlgebraKitMooncakeExt.jl

@@ -227,7 +227,7 @@ end
 
 for (f, f_ne, pb, pf, adj) in (
         (:eig_trunc, :eig_trunc_no_error, :eig_trunc_pullback!, :eig_trunc_pushforward!, :eig_trunc_adjoint),
-        (:eigh_trunc, :eigh_trunc_no_error, :eigh_trunc_pullback!, :eigh_trunc_pushforward, :eigh_trunc_adjoint),
+        (:eigh_trunc, :eigh_trunc_no_error, :eigh_trunc_pullback!, :eigh_trunc_pushforward!, :eigh_trunc_adjoint),
     )
     @eval begin
         @is_primitive Mooncake.DefaultCtx Tuple{typeof($f), Any, MatrixAlgebraKit.AbstractAlgorithm}
@@ -254,7 +254,20 @@ for (f, f_ne, pb, pf, adj) in (
             end
             return output_codual, $adj
         end
-        @is_primitive Mooncake.DefaultCtx Mooncake.ReverseMode Tuple{typeof($f_ne), Any, MatrixAlgebraKit.AbstractAlgorithm}
+        function Mooncake.frule!!(::Dual{typeof($f)}, A_dA::Dual, alg_dalg::Dual)
+            # compute primal
+            A, dA = arrayify(A_dA)
+            alg = Mooncake.primal(alg_dalg)
+            output = $f(A, alg)
+            output_dual = Mooncake.zero_dual(output)
+            dD_ = Mooncake.tangent(output_dual)[1]
+            dV_ = Mooncake.tangent(output_dual)[2]
+            D, dD = arrayify(output[1], dD_)
+            V, dV = arrayify(output[2], dV_)
+            $pf(dA, A, (D, V), (dD, dV))
+            return output_dual
+        end
+        @is_primitive Mooncake.DefaultCtx Tuple{typeof($f_ne), Any, MatrixAlgebraKit.AbstractAlgorithm}
         function Mooncake.rrule!!(::CoDual{typeof($f_ne)}, A_dA::CoDual, alg_dalg::CoDual)
             # compute primal
             A, dA = arrayify(A_dA)
@@ -277,11 +290,11 @@ for (f, f_ne, pb, pf, adj) in (
             end
             return output_codual, $adj
         end
-        function Mooncake.frule!!(::Dual{typeof($f)}, A_dA::Dual, alg_dalg::Dual)
+        function Mooncake.frule!!(::Dual{typeof($f_ne)}, A_dA::Dual, alg_dalg::Dual)
             # compute primal
             A, dA = arrayify(A_dA)
             alg = Mooncake.primal(alg_dalg)
-            output = $f(A, alg)
+            output = $f_ne(A, alg)
             output_dual = Mooncake.zero_dual(output)
             dD_ = Mooncake.tangent(output_dual)[1]
             dV_ = Mooncake.tangent(output_dual)[2]
@@ -478,9 +491,7 @@ end
 
 @is_primitive Mooncake.DefaultCtx Tuple{typeof(svd_trunc), Any, MatrixAlgebraKit.AbstractAlgorithm}
 function Mooncake.rrule!!(::CoDual{typeof(svd_trunc)}, A_dA::CoDual, alg_dalg::CoDual)
-    A_ = Mooncake.primal(A_dA)
-    dA_ = Mooncake.tangent(A_dA)
-    A, dA = arrayify(A_, dA_)
+    A, dA = arrayify(A_dA)
     alg = Mooncake.primal(alg_dalg)
     output = svd_trunc(A, alg)
     # fdata call here is necessary to convert complicated Tangent type (e.g. of a Diagonal
@@ -531,12 +542,10 @@ function Mooncake.frule!!(::Dual{typeof(svd_trunc)}, A_dA::Dual, alg_dalg::Dual)
 end
 
 
-@is_primitive Mooncake.DefaultCtx Mooncake.ReverseMode Tuple{typeof(svd_trunc_no_error), Any, MatrixAlgebraKit.AbstractAlgorithm}
+@is_primitive Mooncake.DefaultCtx Tuple{typeof(svd_trunc_no_error), Any, MatrixAlgebraKit.AbstractAlgorithm}
 function Mooncake.rrule!!(::CoDual{typeof(svd_trunc_no_error)}, A_dA::CoDual, alg_dalg::CoDual)
     # compute primal
-    A_ = Mooncake.primal(A_dA)
-    dA_ = Mooncake.tangent(A_dA)
-    A, dA = arrayify(A_, dA_)
+    A, dA = arrayify(A_dA)
     alg = Mooncake.primal(alg_dalg)
     output = svd_trunc_no_error(A, alg)
     # fdata call here is necessary to convert complicated Tangent type (e.g. of a Diagonal
@@ -559,4 +568,29 @@ function Mooncake.rrule!!(::CoDual{typeof(svd_trunc_no_error)}, A_dA::CoDual, al
     return output_codual, svd_trunc_adjoint
 end
 
+function Mooncake.frule!!(::Dual{typeof(svd_trunc_no_error)}, A_dA::Dual, alg_dalg::Dual)
+    # compute primal
+    A, dA = arrayify(A_dA)
+    alg = Mooncake.primal(alg_dalg)
+    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)
+    output = USVᴴtrunc
+    output_dual = Mooncake.zero_dual(output)
+    Utrunc, Strunc, Vᴴtrunc = output
+    dU_, dS_, dVᴴ_ = Mooncake.tangent(output_dual)
+    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
+
 end

src/pushforwards/eig.jl

@@ -13,3 +13,7 @@ function eig_pushforward!(ΔA, A, DV, ΔDV; kwargs...)
 end
 
 function eig_trunc_pushforward!(ΔA, A, DV, ΔDV; kwargs...) end
+
+function eig_vals_pushforward!(ΔA, A, DV, ΔD; kwargs...)
+    return eig_pushforward!(ΔA, A, DV, ΔD; kwargs...) 
+end

src/pushforwards/eigh.jl

@@ -17,3 +17,5 @@ function eigh_pushforward!(dA, A, DV, dDV; kwargs...)
 end
 
 function eigh_trunc_pushforward!(dA, A, DV, dDV; kwargs...) end
+
+function eigh_vals_pushforward!(dA, A, DV, dDV, ind = Colon(); kwargs...) end

src/pushforwards/svd.jl

@@ -49,12 +49,12 @@ function svd_pushforward!(ΔA, A, USVᴴ, ΔUSVᴴ; rank_atol = default_pullback
         fill!(UÃÃV, 0)
         view(UÃÃV, (1:size(aUAV, 1)), size(aUAV, 1) .+ (1:size(aUAV, 2))) .= aUAV
         view(UÃÃV, size(aUAV, 1) .+ (1:size(aUAV, 2)), 1:size(aUAV, 1)) .= aUAV'
-        rhs = vcat(adjoint(Uperp, ΔA, V), Vᴴperp * ΔA' * U)
+        rhs = vcat(adjoint(Uperp * ΔA * Vᴴ), Vᴴperp * ΔA' * U)
         superKM = -sylvester(UÃÃV, Smat, rhs)
         K̇perp = view(superKM, 1:size(aUAV, 2))
         Ṁperp = view(superKM, (size(aUAV, 2) + 1):(size(aUAV, 1) + size(aUAV, 2)))
         ∂U .+= Uperp * K̇perp
-        ∂V .+= Vperp * Ṁperp
+        ∂V .+= Vᴴperp * Ṁperp
     else
         ImUU = (LinearAlgebra.diagm(ones(eltype(U), m)) - vU * vU')
         ImVV = (LinearAlgebra.diagm(ones(eltype(Vᴴ), n)) - vV * vVᴴ)