Expand set of Mooncake rules

.github/workflows/CI.yml

@@ -30,7 +30,8 @@ jobs:
           - symmetries
           - tensors
           - other
-          - autodiff
+          - mooncake
+          - chainrules
         os:
           - ubuntu-latest
           - macOS-latest
@@ -55,7 +56,8 @@ jobs:
           - symmetries
           - tensors
           - other
-          - autodiff
+          - mooncake
+          - chainrules
         os:
           - ubuntu-latest
           - macOS-latest

Project.toml

@@ -22,8 +22,8 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
-cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
+cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1"
 
 [extensions]
 TensorKitAdaptExt = "Adapt"
@@ -34,6 +34,7 @@ TensorKitMooncakeExt = "Mooncake"
 
 [compat]
 Adapt = "4"
+AllocCheck = "0.2.3"
 Aqua = "0.6, 0.7, 0.8"
 ArgParse = "1.2.0"
 CUDA = "5.9"
@@ -42,10 +43,11 @@ ChainRulesTestUtils = "1"
 Combinatorics = "1"
 FiniteDifferences = "0.12"
 GPUArrays = "11.3.1"
+JET = "0.9, 0.10, 0.11"
 LRUCache = "1.0.2"
 LinearAlgebra = "1"
-MatrixAlgebraKit = "0.6.3"
-Mooncake = "0.4.183"
+MatrixAlgebraKit = "0.6.4"
+Mooncake = "0.5"
 OhMyThreads = "0.8.0"
 Printf = "1"
 Random = "1"
@@ -56,14 +58,15 @@ TensorKitSectors = "0.3.5"
 TensorOperations = "5.1"
 Test = "1"
 TestExtras = "0.2,0.3"
-TupleTools = "1.1"
+TupleTools = "1.5"
 VectorInterface = "0.4.8, 0.5"
 Zygote = "0.7"
 cuTENSOR = "2"
 julia = "1.10"
 
 [extras]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+AllocCheck = "9b6a8646-10ed-4001-bbdc-1d2f46dfbb1a"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 ArgParse = "c7e460c6-2fb9-53a9-8c5b-16f535851c63"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
@@ -72,6 +75,7 @@ ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
+JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
@@ -82,4 +86,4 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1"
 
 [targets]
-test = ["ArgParse", "Adapt", "Aqua", "Combinatorics", "CUDA", "cuTENSOR", "GPUArrays", "LinearAlgebra", "SafeTestsets", "TensorOperations", "Test", "TestExtras", "ChainRulesCore", "ChainRulesTestUtils", "FiniteDifferences", "Zygote", "Mooncake"]
+test = ["ArgParse", "Adapt", "Aqua", "AllocCheck", "Combinatorics", "CUDA", "cuTENSOR", "GPUArrays", "LinearAlgebra", "SafeTestsets", "TensorOperations", "Test", "TestExtras", "ChainRulesCore", "ChainRulesTestUtils", "FiniteDifferences", "Zygote", "Mooncake", "JET"]

ext/TensorKitMooncakeExt/TensorKitMooncakeExt.jl

@@ -1,17 +1,25 @@
 module TensorKitMooncakeExt
 
 using Mooncake
-using Mooncake: @zero_derivative, DefaultCtx, ReverseMode, NoRData, CoDual, arrayify, primal
+using Mooncake: @zero_derivative, @is_primitive,
+    DefaultCtx, MinimalCtx, ReverseMode, NoFData, NoRData, NoTangent,
+    CoDual, Dual, arrayify, primal, tangent, zero_fcodual
 using TensorKit
+import TensorKit as TK
+using VectorInterface
 using TensorOperations: TensorOperations, IndexTuple, Index2Tuple, linearize
 import TensorOperations as TO
-using VectorInterface: One, Zero
+using MatrixAlgebraKit
 using TupleTools
-
+using Random: AbstractRNG
 
 include("utility.jl")
 include("tangent.jl")
 include("linalg.jl")
+include("indexmanipulations.jl")
+include("vectorinterface.jl")
 include("tensoroperations.jl")
+include("planaroperations.jl")
+include("factorizations.jl")
 
 end

ext/TensorKitMooncakeExt/factorizations.jl

@@ -0,0 +1,63 @@
+for f in (:svd_compact, :svd_full)
+    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)
+            alg = primal(alg_dalg)
+
+            USVᴴ = $f(A, primal(alg_dalg))
+            USVᴴ_dUSVᴴ = Mooncake.zero_fcodual(USVᴴ)
+            dUSVᴴ = last.(arrayify.(USVᴴ, tangent(USVᴴ_dUSVᴴ)))
+
+            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 is not guaranteed to actually mutate
+    # so we can simply use the non-mutating version instead and avoid having to worry about
+    # storing copies and restoring state
+    f! = Symbol(f, :!)
+    f!_pullback = Symbol(f!, :_pullback)
+    @eval begin
+        @is_primitive DefaultCtx ReverseMode Tuple{typeof($f!), AbstractTensorMap, Any, MatrixAlgebraKit.AbstractAlgorithm}
+        Mooncake.rrule!!(::CoDual{typeof($f!)}, A_dA::CoDual{<:AbstractTensorMap}, USVᴴ_dUSVᴴ::CoDual, alg_dalg::CoDual) =
+            Mooncake.rrule!!(Mooncake.zero_fcodual($f), A_dA, alg_dalg)
+    end
+end
+
+@is_primitive DefaultCtx ReverseMode Tuple{typeof(svd_trunc), AbstractTensorMap, MatrixAlgebraKit.AbstractAlgorithm}
+function Mooncake.rrule!!(
+        ::CoDual{typeof(svd_trunc)},
+        A_dA::CoDual{<:AbstractTensorMap},
+        alg_dalg::CoDual{<:MatrixAlgebraKit.TruncatedAlgorithm}
+    )
+    A, dA = arrayify(A_dA)
+    alg = primal(alg_dalg)
+
+    USVᴴ = svd_compact(A, alg.alg)
+    USVᴴtrunc, ind = MatrixAlgebraKit.truncate(svd_trunc!, USVᴴ, alg.trunc)
+    ϵ = MatrixAlgebraKit.truncation_error(diagview(USVᴴ[2]), ind)
+
+    USVᴴtrunc_dUSVᴴtrunc = Mooncake.zero_fcodual((USVᴴtrunc..., ϵ))
+    dUSVᴴtrunc = last.(arrayify.(USVᴴtrunc, Base.front(tangent(USVᴴtrunc_dUSVᴴtrunc))))
+
+    function svd_trunc_pullback((_, _, _, dϵ)::Tuple{NoRData, NoRData, NoRData, Real})
+        abs(dϵ) ≤ MatrixAlgebraKit.defaulttol(dϵ) ||
+            @warn "Gradient for `svd_trunc` ignores non-zero tangents for truncation error"
+        MatrixAlgebraKit.svd_pullback!(dA, A, USVᴴ, dUSVᴴtrunc, ind)
+        return ntuple(Returns(NoRData()), 3)
+    end
+
+    return USVᴴtrunc_dUSVᴴtrunc, svd_trunc_pullback
+end
+
+@is_primitive DefaultCtx ReverseMode Tuple{typeof(svd_trunc!), AbstractTensorMap, Any, MatrixAlgebraKit.AbstractAlgorithm}
+Mooncake.rrule!!(::CoDual{typeof(svd_trunc!)}, A_dA::CoDual{<:AbstractTensorMap}, USVᴴ_dUSVᴴ::CoDual, alg_dalg::CoDual) =
+    Mooncake.rrule!!(Mooncake.zero_fcodual(svd_trunc), A_dA, alg_dalg)