Pullback rules for truncation methods (#53)

* prepare pullback rules for truncation

* some updates

* Fix small typos

* fix svd_trunc_pullback and add tests

* fix and test eig pullback, including trunc alternative

* fix rank edge case

* default value of `ind` is untruncated

* update docstrings

* remove some allocations

* mark pullbacks as public

* Bump v0.4.1
[skip ci]

---------

Co-authored-by: Lukas Devos <[email protected]>

Author: Jutho

Project.toml

@@ -1,7 +1,7 @@
 name = "MatrixAlgebraKit"
 uuid = "6c742aac-3347-4629-af66-fc926824e5e4"
 authors = ["Jutho <[email protected]> and contributors"]
-version = "0.4.0"
+version = "0.4.1"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

ext/MatrixAlgebraKitChainRulesCoreExt.jl

@@ -1,7 +1,8 @@
 module MatrixAlgebraKitChainRulesCoreExt
 
 using MatrixAlgebraKit
-using MatrixAlgebraKit: copy_input, TruncatedAlgorithm, zero!
+using MatrixAlgebraKit: copy_input, initialize_output, zero!, diagview,
+    TruncatedAlgorithm, findtruncated, findtruncated_svd
 using ChainRulesCore
 using LinearAlgebra
 
@@ -24,7 +25,7 @@ for qr_f in (:qr_compact, :qr_full)
             QR = $(qr_f!)(Ac, QR, alg)
             function qr_pullback(ΔQR)
                 ΔA = zero(A)
-                MatrixAlgebraKit.qr_compact_pullback!(ΔA, QR, unthunk.(ΔQR))
+                MatrixAlgebraKit.qr_compact_pullback!(ΔA, A, QR, unthunk.(ΔQR))
                 return NoTangent(), ΔA, ZeroTangent(), NoTangent()
             end
             function qr_pullback(::Tuple{ZeroTangent, ZeroTangent}) # is this extra definition useful?
@@ -36,7 +37,7 @@ for qr_f in (:qr_compact, :qr_full)
 end
 function ChainRulesCore.rrule(::typeof(qr_null!), A::AbstractMatrix, N, alg)
     Ac = copy_input(qr_full, A)
-    QR = MatrixAlgebraKit.initialize_output(qr_full!, A, alg)
+    QR = initialize_output(qr_full!, A, alg)
     Q, R = qr_full!(Ac, QR, alg)
     N = copy!(N, view(Q, 1:size(A, 1), (size(A, 2) + 1):size(A, 1)))
     function qr_null_pullback(ΔN)
@@ -45,7 +46,7 @@ function ChainRulesCore.rrule(::typeof(qr_null!), A::AbstractMatrix, N, alg)
         minmn = min(m, n)
         ΔQ = zero!(similar(A, (m, m)))
         view(ΔQ, 1:m, (minmn + 1):m) .= unthunk.(ΔN)
-        MatrixAlgebraKit.qr_compact_pullback!(ΔA, (Q, R), (ΔQ, ZeroTangent()))
+        MatrixAlgebraKit.qr_compact_pullback!(ΔA, A, (Q, R), (ΔQ, ZeroTangent()))
         return NoTangent(), ΔA, ZeroTangent(), NoTangent()
     end
     function qr_null_pullback(::ZeroTangent) # is this extra definition useful?
@@ -62,7 +63,7 @@ for lq_f in (:lq_compact, :lq_full)
             LQ = $(lq_f!)(Ac, LQ, alg)
             function lq_pullback(ΔLQ)
                 ΔA = zero(A)
-                MatrixAlgebraKit.lq_compact_pullback!(ΔA, LQ, unthunk.(ΔLQ))
+                MatrixAlgebraKit.lq_compact_pullback!(ΔA, A, LQ, unthunk.(ΔLQ))
                 return NoTangent(), ΔA, ZeroTangent(), NoTangent()
             end
             function lq_pullback(::Tuple{ZeroTangent, ZeroTangent}) # is this extra definition useful?
@@ -74,7 +75,7 @@ for lq_f in (:lq_compact, :lq_full)
 end
 function ChainRulesCore.rrule(::typeof(lq_null!), A::AbstractMatrix, Nᴴ, alg)
     Ac = copy_input(lq_full, A)
-    LQ = MatrixAlgebraKit.initialize_output(lq_full!, A, alg)
+    LQ = initialize_output(lq_full!, A, alg)
     L, Q = lq_full!(Ac, LQ, alg)
     Nᴴ = copy!(Nᴴ, view(Q, (size(A, 1) + 1):size(A, 2), 1:size(A, 2)))
     function lq_null_pullback(ΔNᴴ)
@@ -83,7 +84,7 @@ function ChainRulesCore.rrule(::typeof(lq_null!), A::AbstractMatrix, Nᴴ, alg)
         minmn = min(m, n)
         ΔQ = zero!(similar(A, (n, n)))
         view(ΔQ, (minmn + 1):n, 1:n) .= unthunk.(ΔNᴴ)
-        MatrixAlgebraKit.lq_compact_pullback!(ΔA, (L, Q), (ZeroTangent(), ΔQ))
+        MatrixAlgebraKit.lq_compact_pullback!(ΔA, A, (L, Q), (ZeroTangent(), ΔQ))
         return NoTangent(), ΔA, ZeroTangent(), NoTangent()
     end
     function lq_null_pullback(::ZeroTangent) # is this extra definition useful?
@@ -95,22 +96,46 @@ end
 for eig in (:eig, :eigh)
     eig_f = Symbol(eig, "_full")
     eig_f! = Symbol(eig_f, "!")
-    eig_f_pb! = Symbol(eig, "_full_pullback!")
+    eig_pb! = Symbol(eig, "_pullback!")
     eig_pb = Symbol(eig, "_pullback")
+    eig_t! = Symbol(eig, "_trunc!")
+    eig_t_pb = Symbol(eig, "_trunc_pullback")
+    _make_eig_t_pb = Symbol("_make_", eig_t_pb)
     @eval begin
         function ChainRulesCore.rrule(::typeof($eig_f!), A::AbstractMatrix, DV, alg)
             Ac = copy_input($eig_f, A)
             DV = $(eig_f!)(Ac, DV, alg)
             function $eig_pb(ΔDV)
                 ΔA = zero(A)
-                MatrixAlgebraKit.$eig_f_pb!(ΔA, DV, unthunk.(ΔDV))
+                MatrixAlgebraKit.$eig_pb!(ΔA, A, DV, unthunk.(ΔDV))
                 return NoTangent(), ΔA, ZeroTangent(), NoTangent()
             end
             function $eig_pb(::Tuple{ZeroTangent, ZeroTangent}) # is this extra definition useful?
                 return NoTangent(), ZeroTangent(), ZeroTangent(), NoTangent()
             end
             return DV, $eig_pb
         end
+        function ChainRulesCore.rrule(
+                ::typeof($eig_t!), A::AbstractMatrix, DV,
+                alg::TruncatedAlgorithm
+            )
+            Ac = copy_input($eig_f, A)
+            D, V = $(eig_f!)(Ac, DV, alg.alg)
+            ind = findtruncated(diagview(D), alg.trunc)
+            return (Diagonal(diagview(D)[ind]), V[:, ind]),
+                $(_make_eig_t_pb)(A, (D, V), ind)
+        end
+        function $(_make_eig_t_pb)(A, DV, ind)
+            function $eig_t_pb(ΔDV)
+                ΔA = zero(A)
+                MatrixAlgebraKit.$eig_pb!(ΔA, A, DV, unthunk.(ΔDV), ind)
+                return NoTangent(), ΔA, ZeroTangent(), NoTangent()
+            end
+            function $eig_t_pb(::Tuple{ZeroTangent, ZeroTangent}) # is this extra definition useful?
+                return NoTangent(), ZeroTangent(), ZeroTangent(), NoTangent()
+            end
+            return $eig_t_pb
+        end
     end
 end
 
@@ -122,7 +147,7 @@ for svd_f in (:svd_compact, :svd_full)
             USVᴴ = $(svd_f!)(Ac, USVᴴ, alg)
             function svd_pullback(ΔUSVᴴ)
                 ΔA = zero(A)
-                MatrixAlgebraKit.svd_compact_pullback!(ΔA, USVᴴ, unthunk.(ΔUSVᴴ))
+                MatrixAlgebraKit.svd_pullback!(ΔA, A, USVᴴ, unthunk.(ΔUSVᴴ))
                 return NoTangent(), ΔA, ZeroTangent(), NoTangent()
             end
             function svd_pullback(::Tuple{ZeroTangent, ZeroTangent, ZeroTangent}) # is this extra definition useful?
@@ -134,27 +159,33 @@ for svd_f in (:svd_compact, :svd_full)
 end
 
 function ChainRulesCore.rrule(
-        ::typeof(svd_trunc!), A::AbstractMatrix, USVᴴ, alg::TruncatedAlgorithm
+        ::typeof(svd_trunc!), A::AbstractMatrix, USVᴴ,
+        alg::TruncatedAlgorithm
     )
-    Ac = MatrixAlgebraKit.copy_input(svd_compact, A)
-    USVᴴ = svd_compact!(Ac, USVᴴ, alg.alg)
+    Ac = copy_input(svd_compact, A)
+    U, S, Vᴴ = svd_compact!(Ac, USVᴴ, alg.alg)
+    ind = findtruncated_svd(diagview(S), alg.trunc)
+    return (U[:, ind], Diagonal(diagview(S)[ind]), Vᴴ[ind, :]),
+        _make_svd_trunc_pullback(A, (U, S, Vᴴ), ind)
+end
+function _make_svd_trunc_pullback(A, USVᴴ, ind)
     function svd_trunc_pullback(ΔUSVᴴ)
         ΔA = zero(A)
-        MatrixAlgebraKit.svd_compact_pullback!(ΔA, USVᴴ, unthunk.(ΔUSVᴴ))
+        MatrixAlgebraKit.svd_pullback!(ΔA, A, USVᴴ, unthunk.(ΔUSVᴴ), ind)
         return NoTangent(), ΔA, ZeroTangent(), NoTangent()
     end
     function svd_trunc_pullback(::Tuple{ZeroTangent, ZeroTangent, ZeroTangent}) # is this extra definition useful?
         return NoTangent(), ZeroTangent(), ZeroTangent(), NoTangent()
     end
-    return MatrixAlgebraKit.truncate!(svd_trunc!, USVᴴ, alg.trunc), svd_trunc_pullback
+    return svd_trunc_pullback
 end
 
 function ChainRulesCore.rrule(::typeof(left_polar!), A::AbstractMatrix, WP, alg)
     Ac = copy_input(left_polar, A)
     WP = left_polar!(Ac, WP, alg)
     function left_polar_pullback(ΔWP)
         ΔA = zero(A)
-        MatrixAlgebraKit.left_polar_pullback!(ΔA, WP, unthunk.(ΔWP))
+        MatrixAlgebraKit.left_polar_pullback!(ΔA, A, WP, unthunk.(ΔWP))
         return NoTangent(), ΔA, ZeroTangent(), NoTangent()
     end
     function left_polar_pullback(::Tuple{ZeroTangent, ZeroTangent}) # is this extra definition useful?
@@ -168,7 +199,7 @@ function ChainRulesCore.rrule(::typeof(right_polar!), A::AbstractMatrix, PWᴴ,
     PWᴴ = right_polar!(Ac, PWᴴ, alg)
     function right_polar_pullback(ΔPWᴴ)
         ΔA = zero(A)
-        MatrixAlgebraKit.right_polar_pullback!(ΔA, PWᴴ, unthunk.(ΔPWᴴ))
+        MatrixAlgebraKit.right_polar_pullback!(ΔA, A, PWᴴ, unthunk.(ΔPWᴴ))
         return NoTangent(), ΔA, ZeroTangent(), NoTangent()
     end
     function right_polar_pullback(::Tuple{ZeroTangent, ZeroTangent}) # is this extra definition useful?

src/MatrixAlgebraKit.jl

@@ -53,6 +53,13 @@ export notrunc, truncrank, trunctol, truncerror, truncfilter
             :TruncationByError, :TruncationIntersection
         )
     )
+    eval(
+        Expr(
+            :public, :qr_compact_pullback!, :lq_compact_pullback!, :left_polar_pullback!,
+            :right_polar_pullback!, :eig_pullback!, :eig_trunc_pullback!, :eigh_pullback!,
+            :eigh_trunc_pullback!, :svd_pullback!, :svd_trunc_pullback!
+        )
+    )
 end
 
 include("common/defaults.jl")

src/pullbacks/eig.jl

@@ -1,7 +1,30 @@
-function eig_full_pullback!(
-        ΔA::AbstractMatrix, DV, ΔDV;
+"""
+    eig_pullback!(
+        ΔA::AbstractMatrix, A, DV, ΔDV, [ind];
+        tol = default_pullback_gaugetol(DV[1]),
+        degeneracy_atol = tol,
+        gauge_atol = tol
+    )
+
+Adds the pullback from the full eigenvalue decomposition of `A` to `ΔA`, given the output
+`DV` of `eig_full` and the cotangent `ΔDV` of `eig_full` or `eig_trunc`.
+
+In particular, it is assumed that `A ≈ V * D * inv(V)` with thus
+`size(A) == size(V) == size(D)` and `D` diagonal. For the cotangents, an arbitrary number of
+eigenvectors or eigenvalues can be missing, i.e. for a matrix `A` of size `(n, n)`, `ΔV` can
+have size `(n, pV)` and `diagview(ΔD)` can have length `pD`. In those cases, additionally
+`ind` is required to specify which eigenvectors or eigenvalues are present in `ΔV` or `ΔD`.
+By default, it is assumed that all eigenvectors and eigenvalues are present.
+
+A warning will be printed if the cotangents are not gauge-invariant, i.e. if the restriction
+of `V' * ΔV` to rows `i` and columns `j` for which `abs(D[i] - D[j]) < degeneracy_atol`, is
+not small compared to `gauge_atol`.
+"""
+function eig_pullback!(
+        ΔA::AbstractMatrix, A, DV, ΔDV, ind = Colon();
         tol::Real = default_pullback_gaugetol(DV[1]),
-        degeneracy_atol::Real = tol, gauge_atol::Real = tol
+        degeneracy_atol::Real = tol,
+        gauge_atol::Real = tol
     )
 
     # Basic size checks and determination
@@ -10,35 +33,125 @@ function eig_full_pullback!(
     ΔDmat, ΔV = ΔDV
     n = LinearAlgebra.checksquare(V)
     n == length(D) || throw(DimensionMismatch())
+    (n, n) == size(ΔA) || throw(DimensionMismatch())
 
     if !iszerotangent(ΔV)
-        VdΔV = V' * ΔV
+        n == size(ΔV, 1) || throw(DimensionMismatch())
+        pV = size(ΔV, 2)
+        VᴴΔV = fill!(similar(V), 0)
+        indV = axes(V, 2)[ind]
+        length(indV) == pV || throw(DimensionMismatch())
+        mul!(view(VᴴΔV, :, indV), V', ΔV)
 
         mask = abs.(transpose(D) .- D) .< degeneracy_atol
-        Δgauge = norm(view(VdΔV, mask), Inf)
+        Δgauge = norm(view(VᴴΔV, mask), Inf)
         Δgauge < gauge_atol ||
             @warn "`eig` cotangents sensitive to gauge choice: (|Δgauge| = $Δgauge)"
 
-        VdΔV .*= conj.(inv_safe.(transpose(D) .- D, degeneracy_atol))
+        VᴴΔV .*= conj.(inv_safe.(transpose(D) .- D, degeneracy_atol))
 
         if !iszerotangent(ΔDmat)
-            diagview(VdΔV) .+= diagview(ΔDmat)
+            ΔDvec = diagview(ΔDmat)
+            pD = length(ΔDvec)
+            indD = axes(D, 1)[ind]
+            length(indD) == pD || throw(DimensionMismatch())
+            view(diagview(VᴴΔV), indD) .+= ΔDvec
         end
-        PΔV = V' \ VdΔV
+        PΔV = V' \ VᴴΔV
         if eltype(ΔA) <: Real
-            ΔAc = mul!(VdΔV, PΔV, V') # recycle VdΔV memory
+            ΔAc = mul!(VᴴΔV, PΔV, V') # recycle VdΔV memory
             ΔA .+= real.(ΔAc)
         else
             ΔA = mul!(ΔA, PΔV, V', 1, 1)
         end
     elseif !iszerotangent(ΔDmat)
-        PΔV = V' \ Diagonal(diagview(ΔDmat))
+        ΔDvec = diagview(ΔDmat)
+        pD = length(ΔDvec)
+        indD = axes(D, 1)[ind]
+        length(indD) == pD || throw(DimensionMismatch())
+        Vp = view(V, :, indD)
+        PΔV = Vp' \ Diagonal(ΔDvec)
         if eltype(ΔA) <: Real
-            ΔAc = PΔV * V'
+            ΔAc = PΔV * Vp'
             ΔA .+= real.(ΔAc)
         else
             ΔA = mul!(ΔA, PΔV, V', 1, 1)
         end
     end
     return ΔA
 end
+
+"""
+    eig_trunc_pullback!(
+        ΔA::AbstractMatrix, ΔDV, A, DV;
+        tol = default_pullback_gaugetol(DV[1]),
+        degeneracy_atol = tol,
+        gauge_atol = tol
+    )
+
+Adds the pullback from the truncated eigenvalue decomposition of `A` to `ΔA`, given the
+output `DV` and the cotangent `ΔDV` of `eig_trunc`.
+
+In particular, it is assumed that `A * V ≈ V * D` with `V` a rectangular matrix of
+eigenvectors and `D` diagonal. For the cotangents, it is assumed that if `ΔV` is not zero,
+then it has the same number of columns as `V`, and if `ΔD` is not zero, then it is a
+diagonal matrix of the same size as `D`.
+
+For this method to work correctly, it is also assumed that the remaining eigenvalues
+(not included in `D`) are (sufficiently) separated from those in `D`.
+
+A warning will be printed if the cotangents are not gauge-invariant, i.e. if the restriction
+of `V' * ΔV` to rows `i` and columns `j` for which `abs(D[i] - D[j]) < degeneracy_atol`, is
+not small compared to `gauge_atol`.
+"""
+function eig_trunc_pullback!(
+        ΔA::AbstractMatrix, A, DV, ΔDV;
+        tol::Real = default_pullback_gaugetol(DV[1]),
+        degeneracy_atol::Real = tol,
+        gauge_atol::Real = tol
+    )
+
+    # Basic size checks and determination
+    Dmat, V = DV
+    D = diagview(Dmat)
+    ΔDmat, ΔV = ΔDV
+    (n, p) = size(V)
+    p == length(D) || throw(DimensionMismatch())
+    (n, n) == size(ΔA) || throw(DimensionMismatch())
+    G = V' * V
+
+    if !iszerotangent(ΔV)
+        (n, p) == size(ΔV) || throw(DimensionMismatch())
+        VᴴΔV = V' * ΔV
+        mask = abs.(transpose(D) .- D) .< degeneracy_atol
+        Δgauge = norm(view(VᴴΔV, mask), Inf)
+        Δgauge < gauge_atol ||
+            @warn "`eig` cotangents sensitive to gauge choice: (|Δgauge| = $Δgauge)"
+
+        ΔVperp = ΔV - V * inv(G) * VᴴΔV
+        VᴴΔV .*= conj.(inv_safe.(transpose(D) .- D, degeneracy_atol))
+    else
+        VᴴΔV = zero(G)
+    end
+
+    if !iszerotangent(ΔDmat)
+        ΔDvec = diagview(ΔDmat)
+        p == length(ΔDvec) || throw(DimensionMismatch())
+        diagview(VᴴΔV) .+= ΔDvec
+    end
+    Z = V' \ VᴴΔV
+
+    # add contribution from orthogonal complement
+    PA = A - (A * V) / V
+    Y = mul!(ΔVperp, PA', Z, 1, 1)
+    X = sylvester(PA', -Dmat', Y)
+    Z .+= X
+
+    if eltype(ΔA) <: Real
+        ΔAc = Z * V'
+        ΔA .+= real.(ΔAc)
+    else
+        ΔA = mul!(ΔA, Z, V', 1, 1)
+    end
+    return ΔA
+end