From 24ebd18c19de8fe7be98d25bf7ba53302ad0cec4 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Tue, 14 Jul 2026 19:27:19 +0200 Subject: [PATCH] Boil down indexmanip rules to add_transform --- ext/TensorKitEnzymeExt/indexmanipulations.jl | 97 +++++-------------- .../braid.jl | 28 ++++-- .../permute.jl | 26 +++-- .../transpose.jl | 17 +++- 4 files changed, 73 insertions(+), 95 deletions(-) diff --git a/ext/TensorKitEnzymeExt/indexmanipulations.jl b/ext/TensorKitEnzymeExt/indexmanipulations.jl index c76f17c70..9e3d4054f 100644 --- a/ext/TensorKitEnzymeExt/indexmanipulations.jl +++ b/ext/TensorKitEnzymeExt/indexmanipulations.jl @@ -1,87 +1,27 @@ -for transform in (:permute, :transpose) - transform! = Symbol(transform, :!) - transform_pb = Symbol(transform, :_pullback_dA) - @eval function EnzymeRules.augmented_primal( - config::EnzymeRules.RevConfigWidth{1}, - func::Const{typeof(TK.$transform!)}, - ::Type{RT}, - C::Annotation{<:AbstractTensorMap}, - A::Annotation{<:AbstractTensorMap}, - p::Const{<:Index2Tuple}, - α::Annotation{<:Number}, - β::Annotation{<:Number}, - ba::Const... - ) where {RT} - C_cache = !isa(β, Const) ? copy(C.val) : nothing - A_cache = EnzymeRules.overwritten(config)[3] ? copy(A.val) : nothing - # if we need to compute Δa, it is faster to allocate an intermediate permuted A - # and store that instead of repeating the permutation in the pullback each time. - # effectively, we replace `add_permute` by `add ∘ permute`. - Ap = if !isa(α, Const) - Ap = $transform(A.val, p.val) - add!(C.val, Ap, α.val, β.val) - Ap - else - bavs = map(a -> a.val, ba) - TK.$transform!(C.val, A.val, p.val, α.val, β.val, bavs...) - nothing - end - cache = (C_cache, A_cache, Ap) - primal = EnzymeRules.needs_primal(config) ? C.val : nothing - shadow = EnzymeRules.needs_shadow(config) ? C.dval : nothing - return EnzymeRules.AugmentedReturn(primal, shadow, cache) - end - @eval function EnzymeRules.reverse( - config::EnzymeRules.RevConfigWidth{1}, - func::Const{typeof(TK.$transform!)}, - ::Type{RT}, - cache, - C::Annotation{<:AbstractTensorMap}, - A::Annotation{<:AbstractTensorMap}, - p::Const{<:Index2Tuple}, - α::Annotation{<:Number}, - β::Annotation{<:Number}, - ba::Const... - ) where {RT} - C_cache, A_cache, Ap = cache - Cval = something(C_cache, C.val) - bavs = map(a -> a.val, ba) - # ΔA - if !isa(A, Const) && !isa(C, Const) - Aval = something(A_cache, A.val) - TK.$transform_pb(A.dval, Aval, C.dval, C.val, p.val, α.val, bavs...) - end - Δα = pullback_dα(α, C, Ap) - Δβ = pullback_dβ(β, C, Cval) - !isa(C, Const) && pullback_dC!(C.dval, β.val) - return nothing, nothing, nothing, Δα, Δβ, map(Returns(nothing), ba)... - end -end - function EnzymeRules.augmented_primal( config::EnzymeRules.RevConfigWidth{1}, - func::Const{typeof(TK.braid!)}, + func::Const{typeof(TK.add_transform!)}, ::Type{RT}, C::Annotation{<:AbstractTensorMap}, A::Annotation{<:AbstractTensorMap}, - p::Const{<:Index2Tuple}, - levels::Const{<:IndexTuple}, + p::Annotation{<:Index2Tuple}, + transformer::Annotation, α::Annotation{<:Number}, β::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) ? copy(C.val) : C.val + A_cache = EnzymeRules.overwritten(config)[3] ? copy(A.val) : A.val # 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`. + bavs = map(a -> a.val, ba) Ap = if !isa(α, Const) - Ap = braid(A.val, p.val, levels.val) + Ap = TK.add_transform!(Enzyme.make_zero(C.val), A.val, p.val, transformer.val, One(), Zero(), bavs...) add!(C.val, Ap, α.val, β.val) Ap else - bavs = map(a -> a.val, ba) - TK.braid!(C.val, A.val, p.val, levels.val, α.val, β.val, bavs...) + TK.add_transform!(C.val, A.val, p.val, transformer.val, α.val, β.val, bavs...) nothing end cache = (C_cache, A_cache, Ap) @@ -91,24 +31,33 @@ function EnzymeRules.augmented_primal( end function EnzymeRules.reverse( config::EnzymeRules.RevConfigWidth{1}, - func::Const{typeof(TK.braid!)}, + func::Const{typeof(TK.add_transform!)}, ::Type{RT}, cache, C::Annotation{<:AbstractTensorMap}, A::Annotation{<:AbstractTensorMap}, - p::Const{<:Index2Tuple}, - levels::Const{<:IndexTuple}, + p::Annotation{<:Index2Tuple}, + transformer::Annotation, α::Annotation{<:Number}, β::Annotation{<:Number}, ba::Const... ) where {RT} C_cache, A_cache, Ap = cache - Cval = something(C_cache, C.val) - Aval = something(A_cache, A.val) + Cval = C_cache + Aval = A_cache bavs = map(a -> a.val, ba) # ΔA if !isa(A, Const) && !isa(C, Const) - TK.braid_pb(A.dval, Aval, C.dval, C.val, p.val, levels.val, α.val, bavs...) + ip = invperm(linearize(p.val)) + pΔA = TO.repartition(ip, numout(Aval)) + TC = VectorInterface.promote_scale(Cval, α.val) + if scalartype(A.dval) <: Real && !(TC <: Real) + ΔAc = TO.tensoralloc_add(TC, C.dval, pΔA, false, Val(false)) + TK.add_transform!(ΔAc, C.dval, pΔA, transformer.val, conj(α.val), Zero(), bavs...) + add!(A.dval, real(ΔAc)) + else + TK.add_transform!(A.dval, C.dval, pΔA, transformer.val, conj(α.val), One(), bavs...) + end end Δαr = pullback_dα(α, C, Ap) Δβr = pullback_dβ(β, C, Cval) diff --git a/test/enzyme-indexmanipulations-transform/braid.jl b/test/enzyme-indexmanipulations-transform/braid.jl index 0d4099d13..d955cf792 100644 --- a/test/enzyme-indexmanipulations-transform/braid.jl +++ b/test/enzyme-indexmanipulations-transform/braid.jl @@ -8,11 +8,13 @@ spacelist = ad_spacelist(fast_tests) eltypes = (Float64, ComplexF64) is_ci = get(ENV, "CI", "false") == "true" -Tαs = is_ci ? (Active,) : (Active, Const) -Tβs = is_ci ? (Active,) : (Active, Const) +rTαs = is_ci ? (Active,) : (Active, Const) +rTβs = is_ci ? (Active,) : (Active, Const) +fTαs = is_ci ? (Duplicated,) : (Duplicated, Const) +fTβs = is_ci ? (Duplicated,) : (Duplicated, Const) if !Sys.iswindows() && VERSION > v"1.11.0-rc" - @timedtestset "Enzyme - Index Manipulations (braid!) $(TensorKit.type_repr(sectortype(eltype(V)))) ($T) Tα $Tα Tβ $Tβ" for V in spacelist, T in eltypes, Tα in Tαs, Tβ in Tβs + @timedtestset "Enzyme - Index Manipulations (braid!) $(TensorKit.type_repr(sectortype(eltype(V)))) ($T)" for V in spacelist, T in eltypes atol = default_tol(T) rtol = default_tol(T) Vstr = TensorKit.type_repr(sectortype(eltype(V))) @@ -24,11 +26,21 @@ if !Sys.iswindows() && VERSION > v"1.11.0-rc" p = randcircshift(numout(A), numin(A)) levels = Tuple(randperm(numind(A))) C = randn!(transpose(A, p)) - EnzymeTestUtils.test_reverse(TensorKit.braid!, Duplicated, (C, Duplicated), (A, Duplicated), (p, Const), (levels, Const), (α, Tα), (β, Tβ); atol, rtol, testset_name = "braid! V $Vstr Tα $Tα Tβ $Tβ") - if !(T <: Real) && !is_ci - EnzymeTestUtils.test_reverse(TensorKit.braid!, Duplicated, (C, Duplicated), (real(A), Duplicated), (p, Const), (levels, Const), (α, Tα), (β, Tβ); atol, rtol, testset_name = "braid! V $Vstr Tα $Tα Tβ $Tβ") - EnzymeTestUtils.test_reverse(TensorKit.braid!, Duplicated, (C, Duplicated), (A, Duplicated), (p, Const), (levels, Const), (real(α), Tα), (β, Tβ); atol, rtol, testset_name = "braid! V $Vstr Tα $Tα Tβ $Tβ") - EnzymeTestUtils.test_reverse(TensorKit.braid!, Duplicated, (C, Duplicated), (A, Duplicated), (p, Const), (levels, Const), (real(α), Tα), (real(β), Tβ); atol, rtol, testset_name = "braid! V $Vstr Tα $Tα Tβ $Tβ") + @testset for Tα in rTαs, Tβ in rTβs + EnzymeTestUtils.test_reverse(TensorKit.braid!, Duplicated, (C, Duplicated), (A, Duplicated), (p, Const), (levels, Const), (α, Tα), (β, Tβ); atol, rtol, testset_name = "braid! V $Vstr Tα $Tα Tβ $Tβ") + if !(T <: Real) && !is_ci + EnzymeTestUtils.test_reverse(TensorKit.braid!, Duplicated, (C, Duplicated), (real(A), Duplicated), (p, Const), (levels, Const), (α, Tα), (β, Tβ); atol, rtol, testset_name = "braid! V $Vstr Tα $Tα Tβ $Tβ") + EnzymeTestUtils.test_reverse(TensorKit.braid!, Duplicated, (C, Duplicated), (A, Duplicated), (p, Const), (levels, Const), (real(α), Tα), (β, Tβ); atol, rtol, testset_name = "braid! V $Vstr Tα $Tα Tβ $Tβ") + EnzymeTestUtils.test_reverse(TensorKit.braid!, Duplicated, (C, Duplicated), (A, Duplicated), (p, Const), (levels, Const), (real(α), Tα), (real(β), Tβ); atol, rtol, testset_name = "braid! V $Vstr Tα $Tα Tβ $Tβ") + end + end + @testset for Tα in fTαs, Tβ in fTβs + EnzymeTestUtils.test_forward(TensorKit.braid!, Duplicated, (C, Duplicated), (A, Duplicated), (p, Const), (levels, Const), (α, Tα), (β, Tβ); atol, rtol, testset_name = "braid! V $Vstr Tα $Tα Tβ $Tβ") + if !(T <: Real) && !is_ci + EnzymeTestUtils.test_forward(TensorKit.braid!, Duplicated, (C, Duplicated), (real(A), Duplicated), (p, Const), (levels, Const), (α, Tα), (β, Tβ); atol, rtol, testset_name = "braid! V $Vstr Tα $Tα Tβ $Tβ") + EnzymeTestUtils.test_forward(TensorKit.braid!, Duplicated, (C, Duplicated), (A, Duplicated), (p, Const), (levels, Const), (real(α), Tα), (β, Tβ); atol, rtol, testset_name = "braid! V $Vstr Tα $Tα Tβ $Tβ") + EnzymeTestUtils.test_forward(TensorKit.braid!, Duplicated, (C, Duplicated), (A, Duplicated), (p, Const), (levels, Const), (real(α), Tα), (real(β), Tβ); atol, rtol, testset_name = "braid! V $Vstr Tα $Tα Tβ $Tβ") + end end end end diff --git a/test/enzyme-indexmanipulations-transform/permute.jl b/test/enzyme-indexmanipulations-transform/permute.jl index d78ec1e41..56f49aa3a 100644 --- a/test/enzyme-indexmanipulations-transform/permute.jl +++ b/test/enzyme-indexmanipulations-transform/permute.jl @@ -8,23 +8,29 @@ spacelist = ad_spacelist(fast_tests) eltypes = (Float64, ComplexF64) is_ci = get(ENV, "CI", "false") == "true" -Tαs = is_ci ? (Active,) : (Active, Const) -Tβs = is_ci ? (Active,) : (Active, Const) +rTαs = is_ci ? (Active,) : (Active, Const) +rTβs = is_ci ? (Active,) : (Active, Const) +fTαs = is_ci ? (Duplicated,) : (Duplicated, Const) +fTβs = is_ci ? (Duplicated,) : (Duplicated, Const) if VERSION >= v"1.11.0-rc" # segfault issues on 1.10 @timedtestset "Enzyme - Index Manipulations (permute!): $(TensorKit.type_repr(sectortype(eltype(V)))) ($T)" for V in spacelist, T in eltypes - println(TensorKit.type_repr(sectortype(eltype(V)))) atol = default_tol(T) rtol = default_tol(T) symmetricbraiding = BraidingStyle(sectortype(eltype(V))) isa SymmetricBraiding + A = randn(T, V[1] ⊗ V[2] ← (V[3] ⊗ V[4] ⊗ V[5])') + α = randn(T) + β = randn(T) + p = randindextuple(numind(A)) + C = randn!(permute(A, p)) - symmetricbraiding && @timedtestset "permute! Tα $Tα, Tβ $Tβ" for Tα in Tαs, Tβ in Tβs - A = randn(T, V[1] ⊗ V[2] ← (V[3] ⊗ V[4] ⊗ V[5])') - α = randn(T) - β = randn(T) - p = randindextuple(numind(A)) - C = randn!(permute(A, p)) - EnzymeTestUtils.test_reverse(TensorKit.permute!, Duplicated, (C, Duplicated), (A, Duplicated), (p, Const), (α, Tα), (β, Tβ); atol, rtol) + symmetricbraiding && @timedtestset "permute!" begin + @testset for Tα in rTαs, Tβ in rTβs + EnzymeTestUtils.test_reverse(TensorKit.permute!, Duplicated, (copy(C), Duplicated), (A, Duplicated), (p, Const), (α, Tα), (β, Tβ); atol, rtol) + end + @testset for Tα in fTαs, Tβ in fTβs + EnzymeTestUtils.test_forward(TensorKit.permute!, Duplicated, (copy(C), Duplicated), (A, Duplicated), (p, Const), (α, Tα), (β, Tβ); atol, rtol) + end end end end diff --git a/test/enzyme-indexmanipulations-transform/transpose.jl b/test/enzyme-indexmanipulations-transform/transpose.jl index 7ba3e2d42..64c72f1b3 100644 --- a/test/enzyme-indexmanipulations-transform/transpose.jl +++ b/test/enzyme-indexmanipulations-transform/transpose.jl @@ -10,8 +10,10 @@ eltypes = (Float64, ComplexF64) is_ci = get(ENV, "CI", "false") == "true" -Tαs = is_ci ? (Active,) : (Active, Const) -Tβs = is_ci ? (Active,) : (Active, Const) +rTαs = is_ci ? (Active,) : (Active, Const) +rTβs = is_ci ? (Active,) : (Active, Const) +fTαs = is_ci ? (Duplicated,) : (Duplicated, Const) +fTβs = is_ci ? (Duplicated,) : (Duplicated, Const) if VERSION > v"1.11.0-rc" # https://github.com/QuantumKitHub/TensorKit.jl/issues/457 @timedtestset "Enzyme - Index Manipulations (transpose!) $(TensorKit.type_repr(sectortype(eltype(V)))) ($T)" for V in spacelist, T in eltypes @@ -25,7 +27,8 @@ if VERSION > v"1.11.0-rc" # https://github.com/QuantumKitHub/TensorKit.jl/issues p = randcircshift(numout(A), numin(A)) C = randn!(transpose(A, p)) !is_ci && EnzymeTestUtils.test_reverse(TensorKit.transpose!, Duplicated, (copy(C), Duplicated), (A, Duplicated), (p, Const), (One(), Const), (Zero(), Const); atol, rtol) - @testset for Tα in Tαs, Tβ in Tβs + !is_ci && EnzymeTestUtils.test_forward(TensorKit.transpose!, Duplicated, (copy(C), Duplicated), (A, Duplicated), (p, Const), (One(), Const), (Zero(), Const); atol, rtol) + @testset for Tα in rTαs, Tβ in rTβs EnzymeTestUtils.test_reverse(TensorKit.transpose!, Duplicated, (copy(C), Duplicated), (A, Duplicated), (p, Const), (α, Tα), (β, Tβ); atol, rtol) if !(T <: Real) && !is_ci EnzymeTestUtils.test_reverse(TensorKit.transpose!, Duplicated, (copy(C), Duplicated), (real(A), Duplicated), (p, Const), (α, Tα), (β, Tβ); atol, rtol) @@ -33,5 +36,13 @@ if VERSION > v"1.11.0-rc" # https://github.com/QuantumKitHub/TensorKit.jl/issues EnzymeTestUtils.test_reverse(TensorKit.transpose!, Duplicated, (copy(C), Duplicated), (real(A), Duplicated), (p, Const), (real(α), Tα), (β, Tβ); atol, rtol) end end + @testset for Tα in fTαs, Tβ in fTβs + EnzymeTestUtils.test_forward(TensorKit.transpose!, Duplicated, (copy(C), Duplicated), (A, Duplicated), (p, Const), (α, Tα), (β, Tβ); atol, rtol) + if !(T <: Real) && !is_ci + EnzymeTestUtils.test_forward(TensorKit.transpose!, Duplicated, (copy(C), Duplicated), (real(A), Duplicated), (p, Const), (α, Tα), (β, Tβ); atol, rtol) + EnzymeTestUtils.test_forward(TensorKit.transpose!, Duplicated, (copy(C), Duplicated), (A, Duplicated), (p, Const), (real(α), Tα), (β, Tβ); atol, rtol) + EnzymeTestUtils.test_forward(TensorKit.transpose!, Duplicated, (copy(C), Duplicated), (real(A), Duplicated), (p, Const), (real(α), Tα), (β, Tβ); atol, rtol) + end + end end end