Reverse Enzyme rules, refactoring, and tests for planar ops

ext/TensorKitEnzymeExt/TensorKitEnzymeExt.jl

@@ -13,5 +13,6 @@ using Random: AbstractRNG
 include("utility.jl")
 include("linalg.jl")
 include("indexmanipulations.jl")
+include("planaroperations.jl")
 
 end

ext/TensorKitEnzymeExt/planaroperations.jl

@@ -0,0 +1,61 @@
+# planartrace!
+# ------------
+# TODO: Fix planartrace pullback
+# This implementation is slightly more involved than its non-planar counterpart
+# this is because we lack a general `pAB` argument in `planarcontract`, and need
+# to keep things planar along the way.
+# In particular, we can't simply tensor product with multiple identities in one go
+# if they aren't "contiguous", e.g. p = ((1, 4, 5), ()), q = ((2, 6), (3, 7))
+
+function EnzymeRules.augmented_primal(
+        config::EnzymeRules.RevConfigWidth{1},
+        ::Const{typeof(TensorKit.planartrace!)},
+        ::Type{RT},
+        C::Annotation{<:AbstractTensorMap},
+        A::Annotation{<:AbstractTensorMap},
+        p::Const{<:Index2Tuple},
+        q::Const{<:Index2Tuple},
+        α::Annotation{<:Number},
+        β::Annotation{<:Number},
+        backend::Const, allocator::Const
+    ) where {RT}
+    cacheC = !isa(β, Const) && copy(C.val)
+    cacheA = EnzymeRules.overwritten(config)[3] ? copy(A.val) : nothing
+
+    TensorKit.planartrace!(C.val, A.val, p.val, q.val, α.val, β.val, backend.val, allocator.val)
+    primal = EnzymeRules.needs_primal(config) ? C.val : nothing
+    shadow = EnzymeRules.needs_shadow(config) ? C.dval : nothing
+    return EnzymeRules.AugmentedReturn(primal, shadow, (cacheC, cacheA))
+end
+function EnzymeRules.reverse(
+        config::EnzymeRules.RevConfigWidth{1},
+        ::Const{typeof(TensorKit.planartrace!)},
+        ::Type{RT},
+        cache,
+        C::Annotation{<:AbstractTensorMap},
+        A::Annotation{<:AbstractTensorMap},
+        p::Const{<:Index2Tuple},
+        q::Const{<:Index2Tuple},
+        α::Annotation{<:Number},
+        β::Annotation{<:Number},
+        backend::Const, allocator::Const
+    ) where {RT}
+    cacheC, cacheA = cache
+    Cval = something(cacheC, C.val)
+    Aval = something(cacheA, A.val)
+    if !isa(C, Const)
+        if !isa(A, Const)
+            TK.planartrace_pullback_ΔA!(A.dval, C.dval, Aval, p.val, q.val, α.val, backend.val, allocator.val)
+        end
+        Δαr = if !isa(α, Const)
+            TK.planartrace_pullback_Δα(C.dval, A.val, p.val, q.val, α.val, backend.val, allocator.val)
+        elseif !isa(α, Const)
+            zero(α.val)
+        else
+            nothing
+        end
+        pullback_dC!(C.dval, β.val)
+    end
+    Δβr = pullback_dβ(β, C, Cval)
+    return nothing, nothing, nothing, nothing, Δαr, Δβr, nothing, nothing
+end

ext/TensorKitMooncakeExt/planaroperations.jl

@@ -41,52 +41,16 @@
 #     function planartrace_pullback(::NoRData)
 #         copy!(C, C_cache)
 #
-#         ΔAr = planartrace_pullback_ΔA!(ΔA, ΔC, A, p, q, α, backend, allocator) # this typically returns NoRData()
-#         Δαr = planartrace_pullback_Δα(ΔC, A, p, q, α, backend, allocator)
+#         ΔAr = TK.planartrace_pullback_ΔA!(ΔA, ΔC, A, p, q, α, backend, allocator) # this typically returns nothing
+#         Δαr = TK.planartrace_pullback_Δα(ΔC, A, p, q, α, backend, allocator)
 #         Δβr = pullback_dβ(ΔC, C, β)
-#         ΔCr = pullback_dC!(ΔC, β) # this typically returns NoRData()
+#         ΔCr = pullback_dC!(ΔC, β) # this typically returns nothing
 #
 #         return NoRData(),
-#             ΔCr, ΔAr, NoRData(), NoRData(),
+#             NoRData(), NoRData(), NoRData(), NoRData(),
 #             Δαr, Δβr, NoRData(), NoRData()
 #     end
 #
 #     return C_ΔC, planartrace_pullback
 # end
 
-# function planartrace_pullback_dA!(
-#         ΔA, ΔC, A, p, q, α, backend, allocator
-#     )
-#     if length(q[1]) == 0
-#         ip = invperm(linearize(p))
-#         pΔA = TK._repartition(ip, A)
-#         TK.transpose!(ΔA, ΔC, pΔA, conj(α), One(), backend, allocator)
-#         return NoRData()
-#     end
-#     # if length(q[1]) == 1
-#     #     ip = invperm((p[1]..., q[2]..., p[2]..., q[1]...))
-#     #     pdA = TK._repartition(ip, A)
-#     #     E = one!(TO.tensoralloc_add(scalartype(A), A, q, false))
-#     #     twist!(E, filter(x -> !isdual(space(E, x)), codomainind(E)))
-#     #     # pE = ((), TK.trivtuple(TO.numind(q)))
-#     #     # pΔC = (TK.trivtuple(TO.numind(p)), ())
-#     #     TensorKit.planaradd!(ΔA, ΔC ⊗ E, pdA, conj(α), One(), backend, allocator)
-#     #     return NoRData()
-#     # end
-#     error("The reverse rule for `planartrace` is not yet implemented")
-# end
-#
-# function planartrace_pullback_dα(
-#         ΔC, A, p, q, α, backend, allocator
-#     )
-#     Tdα = Mooncake.rdata_type(Mooncake.tangent_type(typeof(α)))
-#     Tdα === NoRData && return NoRData()
-#
-#     # TODO: this result might be easier to compute as:
-#     # C′ = βC + α * trace(A) ⟹ At = (C′ - βC) / α
-#     At = TO.tensoralloc_add(scalartype(A), A, p, false, Val(true), allocator)
-#     TensorKit.planartrace!(At, A, p, q, One(), Zero(), backend, allocator)
-#     Δα = project_scalar(α, inner(At, ΔC))
-#     TO.tensorfree!(At, allocator)
-#     return Δα
-# end

src/pullbacks/planaroperations.jl

@@ -0,0 +1,33 @@
+function planartrace_pullback_dA!(
+        ΔA, ΔC, A, p, q, α, backend, allocator
+    )
+    if length(q[1]) == 0
+        ip = invperm(linearize(p))
+        pΔA = TO.repartition(ip, numout(A))
+        transpose!(ΔA, ΔC, pΔA, conj(α), One(), backend, allocator)
+        return nothing
+    end
+    if length(q[1]) == 1
+        ip = invperm((p[1]..., q[2]..., p[2]..., q[1]...))
+        pdA = TO.repartition(ip, numout(A))
+        E = one!(TO.tensoralloc_add(scalartype(A), A, q, false))
+        twist!(E, filter(x -> !isdual(space(E, x)), codomainind(E)))
+        pE = ((), trivtuple(TO.numind(q)))
+        pΔC = (trivtuple(TO.numind(p)), ())
+        planaradd!(ΔA, ΔC ⊗ E, pdA, conj(α), One(), backend, allocator)
+        return nothing
+    end
+    error("The reverse rule for `planartrace` is not yet implemented")
+end
+
+function planartrace_pullback_dα(
+        ΔC, A, p, q, α, backend, allocator
+    )
+    # TODO: this result might be easier to compute as:
+    # C′ = βC + α * trace(A) ⟹ At = (C′ - βC) / α
+    At = TO.tensoralloc_add(scalartype(A), A, p, false, Val(true), allocator)
+    planartrace!(At, A, p, q, One(), Zero(), backend, allocator)
+    Δα = project_scalar(α, inner(At, ΔC))
+    TO.tensorfree!(At, allocator)
+    return Δα
+end

test/enzyme-planaroperations/planarcontract.jl

@@ -0,0 +1,55 @@
+using Test, TestExtras
+using TensorKit
+using TensorOperations
+using VectorInterface: Zero, One
+using Enzyme, EnzymeTestUtils
+using Random
+
+spacelist = ad_spacelist(fast_tests)
+eltypes = (Float64, ComplexF64)
+
+@timedtestset "Enzyme - PlanarOperations (planarcontract): $(TensorKit.type_repr(sectortype(eltype(V)))) ($T)" for V in spacelist, T in eltypes
+    atol = default_tol(T)
+    rtol = default_tol(T)
+    V1, V2, V3, V4, V5 = V
+    k1 = 3
+    k2 = 2
+    k3 = 3
+    k′ = rand(0:(k1 + k2))
+    pA = randcircshift(k′, k1 + k2 - k′, k1)
+    ipA = _repartition(invperm(linearize(pA)), k′)
+    k′ = rand(0:(k2 + k3))
+    pB = randcircshift(k′, k2 + k3 - k′, k2)
+    ipB = _repartition(invperm(linearize(pB)), k′)
+    # TODO: primal value already is broken for this?
+    # pAB = randcircshift(k1, k3)
+    pAB = _repartition(tuple((1:(k1 + k3))...), k1)
+
+    α = randn(T)
+    β = randn(T)
+
+    A = randn(T, permute(V1 ⊗ V2 ⊗ V3 ← (V4 ⊗ V5)', ipA))
+    B = randn(T, permute((V4 ⊗ V5)' ← V1 ⊗ V2 ⊗ V3, ipB))
+    C = randn!(
+        TensorOperations.tensoralloc_contract(
+            T, A, pA, false, B, pB, false, pAB, Val(false)
+        )
+    )
+
+    α = randn(T)
+    β = randn(T)
+
+    A = randn(T, permute(V1 ⊗ V2 ⊗ V3 ← (V4 ⊗ V5)', ipA))
+    B = randn(T, permute((V4 ⊗ V5)' ← V1 ⊗ V2 ⊗ V3, ipB))
+    C = randn!(
+        TensorOperations.tensoralloc_contract(
+            T, A, pA, false, B, pB, false, pAB, Val(false)
+        )
+    )
+    @testset for TC in (Duplicated,), TA in (Duplicated,), TB in (Duplicated,)
+        EnzymeTestUtils.test_reverse(TensorKit.planarcontract!, TC, (C, TC), (A, TA), (pA, Const), (B, TB), (pB, Const), (pAB, Const), (One(), Const), (Zero(), Const); atol, rtol, testset_name = "planarcontract! α = One, β = Zero")
+    end
+    @testset for TC in (Duplicated,), TA in (Duplicated,), TB in (Duplicated,), Tα in (Const, Active), Tβ in (Const, Active)
+        EnzymeTestUtils.test_reverse(TensorKit.planarcontract!, TC, (C, TC), (A, TA), (pA, Const), (B, TB), (pB, Const), (pAB, Const), (α, Tα), (β, Tβ); atol, rtol, testset_name = "planarcontract! Tα = $Tα, Tβ = $Tβ")
+    end
+end

test/enzyme-planaroperations/planartrace.jl

@@ -0,0 +1,37 @@
+using Test, TestExtras
+using TensorKit
+using TensorOperations
+using VectorInterface: Zero, One
+using Enzyme, EnzymeTestUtils
+using Random
+
+spacelist = ad_spacelist(fast_tests)
+eltypes = (Float64, ComplexF64)
+
+@timedtestset "Enzyme - PlanarOperations (planartrace): $(TensorKit.type_repr(sectortype(eltype(V)))) ($T)" for V in spacelist, T in eltypes
+    atol = default_tol(T)
+    rtol = default_tol(T)
+    for _ in 1:5
+        k1 = rand(0:2)
+        k2 = rand(0:1)
+        V1 = map(v -> rand(Bool) ? v' : v, rand(V, k1))
+        V2 = map(v -> rand(Bool) ? v' : v, rand(V, k2))
+        V3 = prod(x -> x ⊗ x', V2[1:k2]; init = one(V[1]))
+        V4 = prod(x -> x ⊗ x', V2[(k2 + 1):end]; init = one(V[1]))
+
+        k′ = rand(0:(k1 + 2k2))
+        (_p, _q) = randcircshift(k′, k1 + 2k2 - k′, k1)
+        p = _repartition(_p, rand(0:k1))
+        q = (tuple(_q[1:2:end]...), tuple(_q[2:2:end]...))
+        ip = _repartition(invperm(linearize((_p, _q))), k′)
+        A = randn(T, permute(prod(V1) ⊗ V3 ← V4, ip))
+
+        α = randn(T)
+        β = randn(T)
+        C = randn!(TensorOperations.tensoralloc_add(T, A, p, false, Val(false)))
+        EnzymeTestUtils.test_reverse(TensorKit.planartrace!, Active, (C, Duplicated), (A, Duplicated), (p, Const), (q, Const), (α, Const), (β, Const), (TensorOperations.DefaultBackend(), Const), (TensorOperations.DefaultAllocator(), Const); atol, rtol)
+        EnzymeTestUtils.test_reverse(TensorKit.planartrace!, Active, (C, Duplicated), (A, Duplicated), (p, Const), (q, Const), (α, Active), (β, Const), (TensorOperations.DefaultBackend(), Const), (TensorOperations.DefaultAllocator(), Const); atol, rtol)
+        EnzymeTestUtils.test_reverse(TensorKit.planartrace!, Active, (C, Duplicated), (A, Duplicated), (p, Const), (q, Const), (α, Const), (β, Active), (TensorOperations.DefaultBackend(), Const), (TensorOperations.DefaultAllocator(), Const); atol, rtol)
+        EnzymeTestUtils.test_reverse(TensorKit.planartrace!, Active, (C, Duplicated), (A, Duplicated), (p, Const), (q, Const), (α, Active), (β, Active), (TensorOperations.DefaultBackend(), Const), (TensorOperations.DefaultAllocator(), Const); atol, rtol)
+    end
+end