add chainrules support

Author: lkdvos

ext/MatrixAlgebraKitChainRulesCoreExt.jl

@@ -188,6 +188,15 @@ for svd_f in (:svd_compact, :svd_full)
             end
             return USVᴴ, svd_pullback
         end
+        function ChainRulesCore.rrule(::typeof($svd_f), A, alg)
+            USVᴴ = $(svd_f)(A, alg)
+            function svd_pullback(ΔUSVᴴ)
+                ΔA = zero(A)
+                MatrixAlgebraKit.svd_pullback!(ΔA, A, USVᴴ, unthunk.(ΔUSVᴴ))
+                return NoTangent(), ΔA, NoTangent()
+            end
+            return USVᴴ, svd_pullback
+        end
     end
 end
 
@@ -198,6 +207,12 @@ function ChainRulesCore.rrule(::typeof(svd_trunc!), A, USVᴴ, alg::TruncatedAlg
     ϵ = truncation_error(diagview(USVᴴ[2]), ind)
     return (USVᴴ′..., ϵ), _make_svd_trunc_pullback(A, USVᴴ, ind)
 end
+function ChainRulesCore.rrule(::typeof(svd_trunc), A, alg::TruncatedAlgorithm)
+    USVᴴ = svd_compact(A, alg.alg)
+    USVᴴ′, ind = MatrixAlgebraKit.truncate(svd_trunc!, USVᴴ, alg.trunc)
+    ϵ = truncation_error(diagview(USVᴴ[2]), ind)
+    return (USVᴴ′..., ϵ), _make_svd_trunc_pullback(A, USVᴴ, ind)
+end
 function _make_svd_trunc_pullback(A, USVᴴ, ind)
     function svd_trunc_pullback(ΔUSVᴴϵ)
         ΔA = zero(A)
@@ -220,6 +235,11 @@ function ChainRulesCore.rrule(::typeof(svd_trunc_no_error!), A, USVᴴ, alg::Tru
     USVᴴ′, ind = MatrixAlgebraKit.truncate(svd_trunc!, USVᴴ, alg.trunc)
     return USVᴴ′, _make_svd_trunc_no_error_pullback(A, USVᴴ, ind)
 end
+function ChainRulesCore.rrule(::typeof(svd_trunc_no_error), A, alg::TruncatedAlgorithm)
+    USVᴴ = svd_compact(A, alg.alg)
+    USVᴴ′, ind = MatrixAlgebraKit.truncate(svd_trunc!, USVᴴ, alg.trunc)
+    return USVᴴ′, _make_svd_trunc_no_error_pullback(A, USVᴴ, ind)
+end
 function _make_svd_trunc_no_error_pullback(A, USVᴴ, ind)
     function svd_trunc_pullback(ΔUSVᴴ)
         ΔA = zero(A)
@@ -240,7 +260,19 @@ function ChainRulesCore.rrule(::typeof(svd_vals!), A, S, alg)
         MatrixAlgebraKit.svd_vals_pullback!(ΔA, A, USVᴴ, unthunk(ΔS))
         return NoTangent(), ΔA, ZeroTangent(), NoTangent()
     end
-    function svd_pullback(::ZeroTangent) # is this extra definition useful?
+    function svd_vals_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(svd_vals), A, 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_vals_pullback(::ZeroTangent) # is this extra definition useful?
         return NoTangent(), ZeroTangent(), ZeroTangent(), NoTangent()
     end
     return diagview(USVᴴ[2]), svd_vals_pullback