Now with tangents

Author: kshyatt

.github/workflows/CI.yml

@@ -31,7 +31,11 @@ jobs:
           - tensors
           - other
           - mooncake
-          - enzyme
+          - enzyme/factorizations
+          - enzyme/linalg
+          - enzyme/tensoroperations
+          - enzyme/vectorinterface
+          - enzyme/indexmanipulations
           - chainrules
         os:
           - ubuntu-latest

Project.toml

@@ -22,6 +22,7 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+EnzymeTestUtils = "12d8515a-0907-448a-8884-5fe00fdf1c5a"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1"
@@ -31,6 +32,7 @@ TensorKitAdaptExt = "Adapt"
 TensorKitCUDAExt = ["CUDA", "cuTENSOR"]
 TensorKitChainRulesCoreExt = "ChainRulesCore"
 TensorKitEnzymeExt = "Enzyme"
+TensorKitEnzymeTestUtilsExt = "EnzymeTestUtils"
 TensorKitFiniteDifferencesExt = "FiniteDifferences"
 TensorKitMooncakeExt = "Mooncake"
 

ext/TensorKitEnzymeExt/factorizations.jl

@@ -28,9 +28,10 @@ for (f, pb) in (
                 alg::Const,
             ) where {RT}
             ret = $f(A.val, alg.val)
-            dret = make_zero(ret)
-            cache = (ret, dret)
-            return EnzymeRules.AugmentedReturn(ret, dret, cache)
+            primal = EnzymeRules.needs_primal(config) ? ret : nothing
+            shadow = EnzymeRules.needs_shadow(config) ? make_zero(ret) : nothing
+            cache = (ret, shadow)
+            return EnzymeRules.AugmentedReturn(primal, shadow, cache)
         end
         function EnzymeRules.reverse(
                 config::EnzymeRules.RevConfigWidth{1},
@@ -40,8 +41,7 @@ for (f, pb) in (
                 A::Annotation{<:AbstractTensorMap},
                 alg::Const,
             ) where {RT}
-            ret, dret = cache
-            $pb(A.dval, A.val, ret, dret)
+            !isa(A, Const) && $pb(A.dval, A.val, cache...)
             return (nothing, nothing)
         end
     end
@@ -57,9 +57,10 @@ for f in (:svd_compact, :svd_full)
                 alg::Const,
             ) where {RT}
             USVᴴ = $f(A.val, alg.val)
-            dUSVᴴ = make_zero(USVᴴ)
-            cache = (USVᴴ, dUSVᴴ)
-            return EnzymeRules.AugmentedReturn(USVᴴ, dUSVᴴ, cache)
+            primal = EnzymeRules.needs_primal(config) ? USVᴴ : nothing
+            shadow = EnzymeRules.needs_shadow(config) ? make_zero(USVᴴ) : nothing
+            cache = (USVᴴ, shadow)
+            return EnzymeRules.AugmentedReturn(primal, shadow, cache)
         end
         function EnzymeRules.reverse(
                 config::EnzymeRules.RevConfigWidth{1},
@@ -69,8 +70,7 @@ for f in (:svd_compact, :svd_full)
                 A::Annotation{<:AbstractTensorMap},
                 alg::Const,
             ) where {RT}
-            USVᴴ, dUSVᴴ = cache
-            MatrixAlgebraKit.svd_pullback!(A.dval, A.val, USVᴴ, dUSVᴴ)
+            !isa(A, Const) && MatrixAlgebraKit.svd_pullback!(A.dval, A.val, cache...)
             return (nothing, nothing)
         end
     end

ext/TensorKitEnzymeExt/indexmanipulations.jl

@@ -90,8 +90,8 @@ function EnzymeRules.augmented_primal(
         β::Annotation{<:Number},
         ba::Const...
     ) where {RT}
-    C_cache = !isa(β, Const) ? copy(C.val) : nothing
-    A_cache = EnzymeRules.overwritten(config)[3] ? copy(A.val) : nothing
+    C_cache = !isa(β, Const) ? deepcopy(C.val) : nothing
+    A_cache = EnzymeRules.overwritten(config)[3] ? deepcopy(A.val) : nothing
     # if we need to compute Δa, it is faster to allocate an intermediate braided A
     # and store that instead of repeating the permutation in the pullback each time.
     # effectively, we replace `add_permute` by `add ∘ permute`.

ext/TensorKitEnzymeExt/linalg.jl

@@ -110,7 +110,47 @@ function EnzymeRules.reverse(
     )
     return (nothing,)
 end
-
+function EnzymeRules.augmented_primal(
+        config::EnzymeRules.RevConfigWidth{1},
+        func::Const{typeof(norm)},
+        ::Type{RT},
+        A::Annotation{<:AbstractTensorMap},
+        p::Const{<:Real},
+    ) where {RT}
+    ret = func.val(A.val, p.val)
+    primal = EnzymeRules.needs_primal(config) ? ret : nothing
+    shadow = EnzymeRules.needs_shadow(config) ? zero(ret) : nothing
+    cacheA = EnzymeRules.overwritten(config)[2] ? copy(A.val) : nothing
+    cache = (ret, cacheA)
+    return EnzymeRules.AugmentedReturn(primal, shadow, cache)
+end
+function EnzymeRules.reverse(
+        config::EnzymeRules.RevConfigWidth{1},
+        func::Const{typeof(norm)},
+        dret::Active,
+        cache,
+        A::Annotation{<:AbstractTensorMap},
+        p::Const{<:Real},
+    )
+    n, cacheA = cache
+    Δn = dret.val
+    Aval = something(cacheA, A.val)
+    if !isa(A, Const)
+        x = (Δn' + Δn) / 2 / hypot(n, eps(one(n)))
+        add!(A.dval, A.val, x)
+    end
+    return (nothing, nothing)
+end
+function EnzymeRules.reverse(
+        config::EnzymeRules.RevConfigWidth{1},
+        func::Const{typeof(norm)},
+        ::Type{<:Const},
+        cache,
+        A::Annotation{<:AbstractTensorMap},
+        p::Const{<:Real},
+    )
+    return (nothing, nothing)
+end
 function EnzymeRules.augmented_primal(
         config::EnzymeRules.RevConfigWidth{1},
         func::Const{typeof(inv)},

ext/TensorKitEnzymeExt/utility.jl

@@ -54,6 +54,9 @@ end
 # Ignore derivatives
 # ------------------
 
+@inline EnzymeRules.inactive_type(::Type{<:TensorKit.FusionTree}) = true
+@inline EnzymeRules.inactive_type(::Type{<:TensorKit.GenericTreeTransformer}) = true
+
 @inline EnzymeRules.inactive(::typeof(TensorKit.fusionblockstructure), arg) = nothing
 @inline EnzymeRules.inactive(::typeof(TensorKit.select), s::HomSpace, i::Index2Tuple) = nothing
 @inline EnzymeRules.inactive(::typeof(TensorKit.flip), s::HomSpace, i::Any) = nothing

ext/TensorKitEnzymeTestUtilsExt.jl

@@ -0,0 +1,60 @@
+module TensorKitEnzymeTestUtilsExt
+
+using TensorKit
+using EnzymeTestUtils
+using EnzymeTestUtils: Enzyme
+import EnzymeTestUtils: to_vec, from_vec, rand_tangent
+
+function EnzymeTestUtils.to_vec(x::TensorMap, seen_vecs::EnzymeTestUtils.AliasDict)
+    has_seen = haskey(seen_vecs, x)
+    is_const = Enzyme.Compiler.guaranteed_const(Core.Typeof(x))
+    if has_seen || is_const
+        x_vec = Float32[]
+    else
+        vec_of_vecs = [b * TensorKit.sqrtdim(c) for (c, b) in blocks(x)]
+        x_vec, back = to_vec(vec_of_vecs)
+        seen_vecs[x] = x_vec
+    end
+    function TensorMap_from_vec(x_vec_new::AbstractVector, seen_xs::EnzymeTestUtils.AliasDict)
+        if xor(has_seen, haskey(seen_xs, x))
+            throw(ErrorException("Arrays must be reconstructed in the same order as they are vectorized."))
+        end
+        has_seen && return seen_xs[x]
+        is_const && return x
+
+        x_new = similar(x)
+        xvec_of_vecs = back(x_vec_new)
+        for (i, (c, b)) in enumerate(blocks(x_new))
+            scale!(b, xvec_of_vecs[i], TensorKit.invsqrtdim(c))
+        end
+        if Core.Typeof(x_new) != Core.Typeof(x)
+            x_new = Core.Typeof(x)(x_new)
+        end
+        seen_xs[x] = x_new
+        return x_new
+    end
+    return x_vec, TensorMap_from_vec
+end
+function EnzymeTestUtils.to_vec(t::TensorKit.AdjointTensorMap, seen_vecs::EnzymeTestUtils.AliasDict)
+    parent_vec, parent_t = to_vec(parent(t), seen_vecs)
+    return parent_vec, adjoint ∘ parent_t
+end
+function EnzymeTestUtils.to_vec(t::TensorKit.DiagonalTensorMap, seen_vecs::EnzymeTestUtils.AliasDict)
+    parent_vec, parent_t = to_vec(TensorMap(t), seen_vecs)
+    return parent_vec, TensorKit.DiagonalTensorMap ∘ parent_t
+end
+
+# generate random tangents for testing
+function EnzymeTestUtils.rand_tangent(rng, t::TensorMap)
+    return TensorMap(EnzymeTestUtils.rand_tangent(rng, t.data), space(t))
+end
+
+function EnzymeTestUtils.rand_tangent(rng, t::TensorKit.AdjointTensorMap)
+    return adjoint(rand_tangent(rng, parent(t)))
+end
+
+function EnzymeTestUtils.rand_tangent(rng, t::DiagonalTensorMap)
+    return DiagonalTensorMap(EnzymeTestUtils.rand_tangent(rng, t.data), space(t, 1))
+end
+
+end

test/enzyme/factorizations/eigh.jl

@@ -6,7 +6,7 @@ using MatrixAlgebraKit
 using Enzyme, EnzymeTestUtils
 using Random
 
-@isdefined(TestSetup) || include("../setup.jl")
+@isdefined(TestSetup) || include("../../setup.jl")
 using .TestSetup
 
 spacelist = (

test/enzyme/factorizations/svd.jl

@@ -6,7 +6,7 @@ using MatrixAlgebraKit
 using Enzyme, EnzymeTestUtils
 using Random
 
-@isdefined(TestSetup) || include("../setup.jl")
+@isdefined(TestSetup) || include("../../setup.jl")
 using .TestSetup
 
 spacelist = (
@@ -58,7 +58,7 @@ end
     atol = default_tol(T)
     rtol = default_tol(T)
     USVᴴ = svd_compact(t)
-    ΔUSVᴴ = (TensorMap(randn!(similar(USVᴴ[1].data)), space(USVᴴ[1])), DiagonalTensorMap(randn!(similar(USVᴴ[2].data)), space(USVᴴ[2], 1)), TensorMap(randn!(similar(USVᴴ[3].data)), space(USVᴴ[3])))
+    ΔUSVᴴ = EnzymeTestUtils.rand_tangent.(USVᴴ)
     remove_svdgauge_dependence!(ΔUSVᴴ[1], ΔUSVᴴ[3], USVᴴ...)
     EnzymeTestUtils.test_reverse(svd_compact, Duplicated, (t, Duplicated); output_tangent = ΔUSVᴴ, atol, rtol)
 
@@ -70,8 +70,8 @@ end
     V_trunc = spacetype(t)(c => min(size(b)...) ÷ 2 for (c, b) in blocks(t))
     trunc = truncspace(V_trunc)
     alg = MatrixAlgebraKit.select_algorithm(svd_trunc_no_error, t, nothing; trunc)
-    USVᴴtrunc = svd_trunc(t, alg)
-    ΔUSVᴴtrunc = randn!(similar.(USVᴴtrunc))
+    USVᴴtrunc = svd_trunc_no_error(t, alg)
+    ΔUSVᴴtrunc = EnzymeTestUtils.rand_tangent.(USVᴴtrunc)
     remove_svdgauge_dependence!(ΔUSVᴴtrunc[1], ΔUSVᴴtrunc[3], USVᴴtrunc...)
     EnzymeTestUtils.test_reverse(svd_trunc_no_error, Duplicated, (t, Duplicated), (alg, Const); output_tangent = ΔUSVᴴtrunc, atol, rtol)
 end

test/enzyme/indexmanipulations/add_braid.jl

@@ -42,23 +42,21 @@ spacelist = (
 eltypes = (Float64, ComplexF64)
 
 @timedtestset "Enzyme - Index Manipulations (add_braid!):" begin
-    @timedtestset "$(TensorKit.type_repr(sectortype(eltype(V)))) ($T)" for V in spacelist, T in eltypes
+    @timedtestset "$(TensorKit.type_repr(sectortype(eltype(V)))) ($T) Tα $Tα Tβ $Tβ" for V in spacelist, T in eltypes, Tα in (Active, Const), Tβ in (Active, Const)
         atol = default_tol(T)
         rtol = default_tol(T)
         Vstr = TensorKit.type_repr(sectortype(eltype(V)))
-        @timedtestset "add_braid! Tα $Tα Tβ $Tβ" for Tα in (Active, Const), Tβ in (Active, Const)
-            A = randn(T, V[1] ⊗ V[2] ← V[4] ⊗ V[5])
-            α = randn(T)
-            β = randn(T)
-            p = randcircshift(numout(A), numin(A))
-            levels = Tuple(randperm(numind(A)))
-            C = randn!(transpose(A, p))
-            EnzymeTestUtils.test_reverse(TensorKit.add_braid!, Duplicated, (C, Duplicated), (A, Duplicated), (p, Const), (levels, Const), (α, Tα), (β, Tβ); atol, rtol, testset_name = "add_braid! V $Vstr Tα $Tα Tβ $Tβ")
-            if !(T <: Real)
-                EnzymeTestUtils.test_reverse(TensorKit.add_braid!, Duplicated, (C, Duplicated), (real(A), Duplicated), (p, Const), (levels, Const), (α, Tα), (β, Tβ); atol, rtol, testset_name = "add_braid! V $Vstr Tα $Tα Tβ $Tβ")
-                EnzymeTestUtils.test_reverse(TensorKit.add_braid!, Duplicated, (C, Duplicated), (A, Duplicated), (p, Const), (levels, Const), (real(α), Tα), (β, Tβ); atol, rtol, testset_name = "add_braid! V $Vstr Tα $Tα Tβ $Tβ")
-                EnzymeTestUtils.test_reverse(TensorKit.add_braid!, Duplicated, (C, Duplicated), (A, Duplicated), (p, Const), (levels, Const), (real(α), Tα), (real(β), Tβ); atol, rtol, testset_name = "add_braid! V $Vstr Tα $Tα Tβ $Tβ")
-            end
+        A = randn(T, V[1] ⊗ V[2] ← V[4] ⊗ V[5])
+        α = randn(T)
+        β = randn(T)
+        p = randcircshift(numout(A), numin(A))
+        levels = Tuple(randperm(numind(A)))
+        C = randn!(transpose(A, p))
+        EnzymeTestUtils.test_reverse(TensorKit.add_braid!, Duplicated, (C, Duplicated), (A, Duplicated), (p, Const), (levels, Const), (α, Tα), (β, Tβ); atol, rtol, testset_name = "add_braid! V $Vstr Tα $Tα Tβ $Tβ")
+        if !(T <: Real)
+            EnzymeTestUtils.test_reverse(TensorKit.add_braid!, Duplicated, (C, Duplicated), (real(A), Duplicated), (p, Const), (levels, Const), (α, Tα), (β, Tβ); atol, rtol, testset_name = "add_braid! V $Vstr Tα $Tα Tβ $Tβ")
+            EnzymeTestUtils.test_reverse(TensorKit.add_braid!, Duplicated, (C, Duplicated), (A, Duplicated), (p, Const), (levels, Const), (real(α), Tα), (β, Tβ); atol, rtol, testset_name = "add_braid! V $Vstr Tα $Tα Tβ $Tβ")
+            EnzymeTestUtils.test_reverse(TensorKit.add_braid!, Duplicated, (C, Duplicated), (A, Duplicated), (p, Const), (levels, Const), (real(α), Tα), (real(β), Tβ); atol, rtol, testset_name = "add_braid! V $Vstr Tα $Tα Tβ $Tβ")
         end
     end
 end