Skip to content

Boil down indexmanip rules to add_transform #486

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
kshyatt wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Boil down indexmanip rules to add_transform #486

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
97 changes: 23 additions & 74 deletions ext/TensorKitEnzymeExt/indexmanipulations.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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)
Expand All @@ -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...)

@lkdvos lkdvos Jul 14, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm really surprised this works though - the transformer is the metadata needed to efficiently permute and recouple, but this depends on the value of p so unless this is accidentally a self-inverse permutation where C and A have the same structure, this should probably give wrong results?

@kshyatt kshyatt Jul 15, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was testing this on spaces that are "too easy" maybe, for the A4 ones it doesn't work :(

@lkdvos lkdvos Jul 15, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, the problem is that at this point we no longer have access to what kind of transformer it is, so it is not possible to construct the reverse one from that.

@kshyatt kshyatt Jul 15, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

But my negative lines diff!!!!!!

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)
Expand Down
28 changes: 20 additions & 8 deletions test/enzyme-indexmanipulations-transform/braid.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)))
Expand All @@ -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
Expand Down
26 changes: 16 additions & 10 deletions test/enzyme-indexmanipulations-transform/permute.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
17 changes: 14 additions & 3 deletions test/enzyme-indexmanipulations-transform/transpose.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -25,13 +27,22 @@ 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)
EnzymeTestUtils.test_reverse(TensorKit.transpose!, Duplicated, (copy(C), Duplicated), (A, Duplicated), (p, Const), (real(α), Tα), (β, Tβ); atol, rtol)
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
Loading