From 0bddff42fbc74f8ec6392282b43a2aca935e285f Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Mon, 6 Jul 2026 14:25:43 -0400 Subject: [PATCH] Reverse Enzyme rules, refactoring, and tests for planar ops --- ext/TensorKitEnzymeExt/TensorKitEnzymeExt.jl | 1 + ext/TensorKitEnzymeExt/planaroperations.jl | 61 +++++++++++++++++++ ext/TensorKitMooncakeExt/planaroperations.jl | 44 ++----------- src/pullbacks/planaroperations.jl | 33 ++++++++++ .../enzyme-planaroperations/planarcontract.jl | 55 +++++++++++++++++ test/enzyme-planaroperations/planartrace.jl | 37 +++++++++++ 6 files changed, 191 insertions(+), 40 deletions(-) create mode 100644 ext/TensorKitEnzymeExt/planaroperations.jl create mode 100644 src/pullbacks/planaroperations.jl create mode 100644 test/enzyme-planaroperations/planarcontract.jl create mode 100644 test/enzyme-planaroperations/planartrace.jl diff --git a/ext/TensorKitEnzymeExt/TensorKitEnzymeExt.jl b/ext/TensorKitEnzymeExt/TensorKitEnzymeExt.jl index c7bde1e2b..b0651cd4b 100644 --- a/ext/TensorKitEnzymeExt/TensorKitEnzymeExt.jl +++ b/ext/TensorKitEnzymeExt/TensorKitEnzymeExt.jl @@ -13,5 +13,6 @@ using Random: AbstractRNG include("utility.jl") include("linalg.jl") include("indexmanipulations.jl") +include("planaroperations.jl") end diff --git a/ext/TensorKitEnzymeExt/planaroperations.jl b/ext/TensorKitEnzymeExt/planaroperations.jl new file mode 100644 index 000000000..f9f21846f --- /dev/null +++ b/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 diff --git a/ext/TensorKitMooncakeExt/planaroperations.jl b/ext/TensorKitMooncakeExt/planaroperations.jl index 48d21a78f..cf9622713 100644 --- a/ext/TensorKitMooncakeExt/planaroperations.jl +++ b/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 diff --git a/src/pullbacks/planaroperations.jl b/src/pullbacks/planaroperations.jl new file mode 100644 index 000000000..a329b1c59 --- /dev/null +++ b/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 diff --git a/test/enzyme-planaroperations/planarcontract.jl b/test/enzyme-planaroperations/planarcontract.jl new file mode 100644 index 000000000..0f164497c --- /dev/null +++ b/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 diff --git a/test/enzyme-planaroperations/planartrace.jl b/test/enzyme-planaroperations/planartrace.jl new file mode 100644 index 000000000..b7d585f60 --- /dev/null +++ b/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