add pullback1, rrule and test for eig(h)_vals`

Author: lkdvos

ext/MatrixAlgebraKitChainRulesCoreExt.jl

@@ -95,6 +95,11 @@ for eig in (:eig, :eigh)
     eig_t! = Symbol(eig, "_trunc!")
     eig_t_pb = Symbol(eig, "_trunc_pullback")
     _make_eig_t_pb = Symbol("_make_", eig_t_pb)
+    eig_v = Symbol(eig, "_vals")
+    eig_v! = Symbol(eig_v, "!")
+    eig_v_pb = Symbol(eig_v, "_pullback")
+    eig_v_pb! = Symbol(eig_v_pb, "!")
+
     @eval begin
         function ChainRulesCore.rrule(::typeof($eig_f!), A, DV, alg)
             Ac = copy_input($eig_f, A)
@@ -131,6 +136,18 @@ for eig in (:eig, :eigh)
             end
             return $eig_t_pb
         end
+        function ChainRulesCore.rrule(::typeof($eig_v!), A, D, alg)
+            DV = $eig_f(A, alg)
+            function $eig_v_pb(ΔD)
+                ΔA = zero(A)
+                MatrixAlgebraKit.$eig_v_pb!(ΔA, A, DV, unthunk(ΔD))
+                return NoTangent(), ΔA, ZeroTangent(), NoTangent()
+            end
+            function $eig_v_pb(::ZeroTangent) # is this extra definition useful?
+                return NoTangent(), ZeroTangent(), ZeroTangent(), NoTangent()
+            end
+            return diagview(DV[1]), $eig_v_pb
+        end
     end
 end
 

src/pullbacks/eig.jl

@@ -151,3 +151,27 @@ function eig_trunc_pullback!(
     end
     return ΔA
 end
+
+"""
+    eig_vals_pullback!(
+        ΔA, A, DV, ΔD, [ind];
+        degeneracy_atol::Real = default_pullback_rank_atol(DV[1]),
+    )
+
+Adds the pullback from the eigenvalues of `A` to `ΔA`, given the output
+`DV` of `eig_full` and the cotangent `ΔD` of `eig_vals`.
+
+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 eigenvalues can be missing, i.e.
+for a matrix `A` of size `(n, n)`, `diagview(ΔD)` can have length `pD`. In those cases,
+additionally `ind` is required to specify which eigenvalues are present in `ΔV` or `ΔD`.
+By default, it is assumed that all eigenvectors and eigenvalues are present.
+"""
+function eig_vals_pullback!(
+        ΔA, A, DV, ΔD, ind = Colon();
+        degeneracy_atol::Real = default_pullback_rank_atol(DV[1]),
+    )
+
+    ΔDV = (diagonal(ΔD), nothing)
+    return eig_pullback!(ΔA, A, DV, ΔDV, ind; degeneracy_atol)
+end

src/pullbacks/eigh.jl

@@ -141,3 +141,27 @@ function eigh_trunc_pullback!(
     end
     return ΔA
 end
+
+"""
+    eigh_vals_pullback!(
+        ΔA, A, DV, ΔD, [ind];
+        degeneracy_atol::Real = default_pullback_rank_atol(DV[1]),
+    )
+
+Adds the pullback from the eigenvalues of `A` to `ΔA`, given the output
+`DV` of `eigh_full` and the cotangent `ΔD` of `eig_vals`.
+
+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 eigenvalues can be missing, i.e.
+for a matrix `A` of size `(n, n)`, `diagview(ΔD)` can have length `pD`. In those cases,
+additionally `ind` is required to specify which eigenvalues are present in `ΔV` or `ΔD`.
+By default, it is assumed that all eigenvectors and eigenvalues are present.
+"""
+function eigh_vals_pullback!(
+        ΔA, A, DV, ΔD, ind = Colon();
+        degeneracy_atol::Real = default_pullback_rank_atol(DV[1]),
+    )
+
+    ΔDV = (diagonal(ΔD), nothing)
+    return eigh_pullback!(ΔA, A, DV, ΔDV, ind; degeneracy_atol)
+end

test/chainrules.jl

@@ -11,22 +11,23 @@ include("ad_utils.jl")
 for f in
     (
         :qr_compact, :qr_full, :qr_null, :lq_compact, :lq_full, :lq_null,
-        :eig_full, :eig_trunc, :eigh_full, :eigh_trunc,
+        :eig_full, :eig_trunc, :eig_vals, :eigh_full, :eigh_trunc, :eigh_vals,
         :svd_compact, :svd_trunc, :svd_vals,
         :left_polar, :right_polar,
     )
     copy_f = Symbol(:copy_, f)
     f! = Symbol(f, '!')
+    _hermitian = startswith(string(f), "eigh")
     @eval begin
         function $copy_f(input, alg)
-            if $f === eigh_full || $f === eigh_trunc
+            if $_hermitian
                 input = (input + input') / 2
             end
             return $f(input, alg)
         end
         function ChainRulesCore.rrule(::typeof($copy_f), input, alg)
             output = MatrixAlgebraKit.initialize_output($f!, input, alg)
-            if $f === eigh_full || $f === eigh_trunc
+            if $_hermitian
                 input = (input + input') / 2
             else
                 input = copy(input)
@@ -229,12 +230,13 @@ end
     ΔD2 = Diagonal(randn(rng, complex(T), m))
     for alg in (LAPACK_Simple(), LAPACK_Expert())
         test_rrule(
-            copy_eig_full, A, alg ⊢ NoTangent();
-            output_tangent = (ΔD, ΔV), atol = atol, rtol = rtol
+            copy_eig_full, A, alg ⊢ NoTangent(); output_tangent = (ΔD, ΔV), atol, rtol
         )
         test_rrule(
-            copy_eig_full, A, alg ⊢ NoTangent();
-            output_tangent = (ΔD2, ΔV), atol = atol, rtol = rtol
+            copy_eig_full, A, alg ⊢ NoTangent(); output_tangent = (ΔD2, ΔV), atol, rtol
+        )
+        test_rrule(
+            copy_eig_vals, A, alg ⊢ NoTangent(); output_tangent = diagview(ΔD), atol, rtol
         )
         for r in 1:4:m
             truncalg = TruncatedAlgorithm(alg, truncrank(r; by = abs))
@@ -305,12 +307,13 @@ end
         )
         # copy_eigh_full includes a projector onto the Hermitian part of the matrix
         test_rrule(
-            copy_eigh_full, A, alg ⊢ NoTangent(); output_tangent = (ΔD, ΔV),
-            atol = atol, rtol = rtol
+            copy_eigh_full, A, alg ⊢ NoTangent(); output_tangent = (ΔD, ΔV), atol, rtol
         )
         test_rrule(
-            copy_eigh_full, A, alg ⊢ NoTangent(); output_tangent = (ΔD2, ΔV),
-            atol = atol, rtol = rtol
+            copy_eigh_full, A, alg ⊢ NoTangent(); output_tangent = (ΔD2, ΔV), atol, rtol
+        )
+        test_rrule(
+            copy_eigh_vals, A, alg ⊢ NoTangent(); output_tangent = diagview(ΔD), atol, rtol
         )
         for r in 1:4:m
             truncalg = TruncatedAlgorithm(alg, truncrank(r; by = abs))