specialize SVD pullback implementations

Author: lkdvos

ext/TensorKitMooncakeExt/TensorKitMooncakeExt.jl

@@ -9,6 +9,7 @@ import TensorKit as TK
 using VectorInterface
 using TensorOperations: TensorOperations, IndexTuple, Index2Tuple, linearize
 import TensorOperations as TO
+using MatrixAlgebraKit
 using TupleTools
 using Random: AbstractRNG
 
@@ -19,5 +20,6 @@ include("indexmanipulations.jl")
 include("vectorinterface.jl")
 include("tensoroperations.jl")
 include("planaroperations.jl")
+include("factorizations.jl")
 
 end

ext/TensorKitMooncakeExt/factorizations.jl

@@ -0,0 +1,67 @@
+for f in (:svd_compact, :svd_full)
+    # non-mutating version
+    # --------------------
+    f_pullback = Symbol(f, :_pullback)
+    @eval begin
+        @is_primitive DefaultCtx ReverseMode Tuple{typeof($f), AbstractTensorMap, MatrixAlgebraKit.AbstractAlgorithm}
+        function Mooncake.rrule!!(::CoDual{typeof($f)}, A_dA::CoDual{<:AbstractTensorMap}, alg_dalg::CoDual)
+            A, dA = arrayify(A_dA)
+            USVᴴ = $f(A, primal(alg_dalg))
+            USVᴴ_dUSVᴴ = Mooncake.zero_fcodual(USVᴴ)
+            (U, dU), (S, dS), (Vᴴ, dVᴴ) = arrayify.(USVᴴ, tangent(USVᴴ_dUSVᴴ))
+            dUSVᴴ = dU, dS, dVᴴ
+
+            function $f_pullback(::NoRData)
+                MatrixAlgebraKit.svd_pullback!(dA, A, USVᴴ, dUSVᴴ)
+                MatrixAlgebraKit.zero!.(dUSVᴴ)
+                return ntuple(Returns(NoRData()), 3)
+            end
+
+            return USVᴴ_dUSVᴴ, $f_pullback
+        end
+    end
+
+    # mutating version
+    # ----------------
+    f! = Symbol(f, :!)
+    f!_pullback = Symbol(f!, :_pullback)
+    @eval begin
+        @is_primitive(
+            DefaultCtx,
+            ReverseMode,
+            Tuple{typeof($f!), AbstractTensorMap, Any, MatrixAlgebraKit.AbstractAlgorithm}
+        )
+
+        function Mooncake.rrule!!(
+                ::CoDual{typeof($f!)},
+                A_dA::CoDual{<:AbstractTensorMap},
+                USVᴴ_dUSVᴴ::CoDual,
+                alg_dalg::CoDual
+            )
+            # unpack values
+            A, dA = arrayify(A_dA)
+            (U, dU), (S, dS), (Vᴴ, dVᴴ) = arrayify.(USVᴴ, tangent(USVᴴ_dUSVᴴ))
+            dUSVᴴ = dU, dS, dVᴴ
+
+            Ac = copy(A)
+            USVᴴc = copy.(USVᴴ)
+
+            output = $f!(A, USVᴴ, primal(alg_dalg))
+            @assert output === USVᴴ "expected in-place algorithm"
+
+            function $f!_pullback(::NoRData)
+                # compute pullbacks
+                MatrixAlgebraKit.svd_pullback!(dA, Ac, USVᴴ, dUSVᴴ)
+                MatrixAlgebraKit.zero!.(dUSVᴴ)
+
+                # restore state
+                copy!(A, Ac)
+                copy!.(USVᴴ, USVᴴc)
+
+                return ntuple(Returns(NoRData()), 4)
+            end
+
+            return USVᴴ_dUSVᴴ, $f!_pullback
+        end
+    end
+end