add pullback, rrule and test for svd_vals

Author: lkdvos

ext/MatrixAlgebraKitChainRulesCoreExt.jl

@@ -176,6 +176,19 @@ function _make_svd_trunc_pullback(A, USVᴴ, ind)
     return svd_trunc_pullback
 end
 
+function ChainRulesCore.rrule(::typeof(svd_vals!), A, S, alg)
+    USVᴴ = svd_compact(A, alg)
+    function svd_vals_pullback(ΔS)
+        ΔA = zero(A)
+        MatrixAlgebraKit.svd_vals_pullback!(ΔA, A, USVᴴ, unthunk(ΔS))
+        return NoTangent(), ΔA, ZeroTangent(), NoTangent()
+    end
+    function svd_pullback(::ZeroTangent) # is this extra definition useful?
+        return NoTangent(), ZeroTangent(), ZeroTangent(), NoTangent()
+    end
+    return diagview(USVᴴ[2]), svd_vals_pullback
+end
+
 function ChainRulesCore.rrule(::typeof(left_polar!), A, WP, alg)
     Ac = copy_input(left_polar, A)
     WP = left_polar!(Ac, WP, alg)

src/pullbacks/svd.jl

@@ -6,7 +6,7 @@
         gauge_atol::Real = default_pullback_gauge_atol(ΔUSVᴴ[1], ΔUSVᴴ[3])
     )
 
-Adds the pullback from the SVD of `A` to `ΔA` given the output USVᴴ of `svd_compact` or
+Adds the pullback from the SVD of `A` to `ΔA` given the output `USVᴴ` of `svd_compact` or
 `svd_full` and the cotangent `ΔUSVᴴ` of `svd_compact`, `svd_full` or `svd_trunc`.
 
 In particular, it is assumed that `A ≈ U * S * Vᴴ`, or thus, that no singular values with
@@ -201,3 +201,29 @@ function svd_trunc_pullback!(
     ΔA = mul!(ΔA, U, Y' * Ṽᴴ, 1, 1)
     return ΔA
 end
+
+"""
+    svd_vals_pullback!(
+        ΔA, A, USVᴴ, ΔS, [ind];
+        rank_atol::Real = default_pullback_rank_atol(USVᴴ[2]),
+        degeneracy_atol::Real = default_pullback_rank_atol(USVᴴ[2])
+    )
+
+
+Adds the pullback from the singular values of `A` to `ΔA`, given the output
+`USVᴴ` of `svd_compact`, and the cotangent `ΔS` of `svd_vals`.
+
+In particular, it is assumed that `A ≈ U * S * Vᴴ`, or thus, that no singular values with
+magnitude less than `rank_atol` are missing from `S`. For the cotangents, an arbitrary
+number of singular vectors or singular values can be missing, i.e. for a matrix `A` with
+size `(m, n)`, `diagview(ΔS)` can have length `pS`. In those cases, additionally `ind` is required to
+specify which singular vectors and values are present in `ΔS`.
+"""
+function svd_vals_pullback!(
+        ΔA, A, USVᴴ, ΔS, ind = Colon();
+        rank_atol::Real = default_pullback_rank_atol(USVᴴ[2]),
+        degeneracy_atol::Real = default_pullback_rank_atol(USVᴴ[2])
+    )
+    ΔUSVᴴ = (nothing, diagonal(ΔS), nothing)
+    return svd_pullback!(ΔA, A, USVᴴ, ΔUSVᴴ, ind; rank_atol, degeneracy_atol)
+end

test/chainrules.jl

@@ -11,7 +11,8 @@ 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, :svd_compact, :svd_trunc,
+        :eig_full, :eig_trunc, :eigh_full, :eigh_trunc,
+        :svd_compact, :svd_trunc, :svd_vals,
         :left_polar, :right_polar,
     )
     copy_f = Symbol(:copy_, f)
@@ -404,6 +405,10 @@ end
                 copy_svd_compact, A, alg ⊢ NoTangent();
                 output_tangent = (ΔU, ΔS2, ΔVᴴ), atol = atol, rtol = rtol
             )
+            test_rrule(
+                copy_svd_vals, A, alg ⊢ NoTangent();
+                output_tangent = diagview(ΔS), atol, rtol
+            )
             for r in 1:4:minmn
                 truncalg = TruncatedAlgorithm(alg, truncrank(r))
                 ind = MatrixAlgebraKit.findtruncated(diagview(S), truncalg.trunc)