Skip to content

Try letting Enzyme make rules for insert/remove unit #488

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

Merged
lkdvos merged 3 commits into from
Jul 15, 2026
+91 −113
Merged

Try letting Enzyme make rules for insert/remove unit #488

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
c157be0
Try letting Enzyme make rules for insert/remove unit
kshyatt Jul 15, 2026
a411580
Run fewer insert tests
kshyatt Jul 15, 2026
6401041
Get rid of unneeded copies
kshyatt Jul 15, 2026
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
2 changes: 1 addition & 1 deletion ext/TensorKitEnzymeExt/TensorKitEnzymeExt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ module TensorKitEnzymeExt
using Enzyme
using TensorKit
import TensorKit as TK
using TensorKit: subblock
using TensorKit: subblock, block
using VectorInterface
using TensorOperations: TensorOperations, IndexTuple, Index2Tuple, linearize
import TensorOperations as TO
Expand Down
89 changes: 0 additions & 89 deletions ext/TensorKitEnzymeExt/indexmanipulations.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -148,92 +148,3 @@ function EnzymeRules.reverse(
end
return (nothing, nothing)
end

for insertunit in (:insertleftunit, :insertrightunit)
@eval begin
function EnzymeRules.augmented_primal(
config::EnzymeRules.RevConfigWidth{1},
func::Const{typeof($insertunit)},
::Type{RT},
tsrc::Annotation{<:AbstractTensorMap},
ival::Const{<:Val};
kwargs...
) where {RT}
if tsrc.val isa TensorMap && !get(kwargs, :copy, false) && !isa(tsrc, Const)
tsrc_cache = tsrc.val
tdst = $insertunit(tsrc.val, ival.val; kwargs...)
Δtdst = $insertunit(tsrc.dval, ival.val; kwargs...)
else
tsrc_cache = nothing
tdst = $insertunit(tsrc.val, ival.val; kwargs...)
Δtdst = make_zero(tdst)
end
primal = EnzymeRules.needs_primal(config) ? tdst : nothing
shadow = EnzymeRules.needs_shadow(config) ? Δtdst : nothing
cache = (tsrc_cache, tdst, Δtdst)
return EnzymeRules.AugmentedReturn(primal, shadow, cache)
end
function EnzymeRules.reverse(
config::EnzymeRules.RevConfigWidth{1},
func::Const{typeof($insertunit)},
::Type{RT},
cache,
tsrc::Annotation{<:AbstractTensorMap},
ival::Const{<:Val};
kwargs...
) where {RT}
tsrc_cache, tdst, Δtdst = cache
# note: since data is already shared for <:TensorMap, don't have to do anything here!
if isnothing(tsrc_cache) && !isa(tsrc, Const)
for (c, b) in blocks(Δtdst)
add!(block(tsrc.dval, c), b)
end
end
return (nothing, nothing)
end
end
end

function EnzymeRules.augmented_primal(
config::EnzymeRules.RevConfigWidth{1},
func::Const{typeof(removeunit)},
::Type{RT},
tsrc::Annotation{<:AbstractTensorMap},
ival::Const{<:Val};
kwargs...
) where {RT}
# tdst shares data with tsrc if <:TensorMap & copy=false, in this case we have to deal with correctly
# sharing address spaces
if tsrc.val isa TensorMap && !get(kwargs, :copy, false) && !isa(tsrc, Const)
tsrc_cache = tsrc.val
tdst = removeunit(tsrc.val, ival.val; kwargs...)
Δtdst = removeunit(tsrc.dval, ival.val)
else
tsrc_cache = nothing
tdst = removeunit(tsrc.val, ival.val; kwargs...)
Δtdst = make_zero(tdst)
end
primal = EnzymeRules.needs_primal(config) ? tdst : nothing
shadow = EnzymeRules.needs_shadow(config) ? Δtdst : nothing
cache = (tsrc_cache, tdst, Δtdst)
return EnzymeRules.AugmentedReturn(primal, shadow, cache)
end

function EnzymeRules.reverse(
config::EnzymeRules.RevConfigWidth{1},
func::Const{typeof(removeunit)},
::Type{RT},
cache,
tsrc::Annotation{<:AbstractTensorMap},
ival::Const{<:Val};
kwargs...
) where {RT}
tsrc_cache, tdst, Δtdst = cache
# note: since data for <: TensorMap is already shared, don't have to do anything here!
if isnothing(tsrc_cache) && !isa(tsrc, Const)
for (c, b) in blocks(Δtdst)
add!(block(tsrc.dval, c), b)
end
end
return (nothing, nothing)
end
60 changes: 57 additions & 3 deletions ext/TensorKitEnzymeExt/utility.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,10 +45,35 @@ function EnzymeRules.reverse(
t::Annotation{<:AbstractTensorMap},
f::Annotation,
) where {RT}
dret = cache
if !isnothing(dret) && !isa(t, Const)
subblock(t.dval, f.val) .= dret
return (nothing, nothing)
end

function EnzymeRules.augmented_primal(
config::EnzymeRules.RevConfigWidth{1},
::Const{typeof(block)},
::Type{RT},
t::Annotation{<:AbstractTensorMap},
c::Annotation{<:Sector},
) where {RT}
ret = EnzymeRules.needs_primal(config) ? block(t.val, c.val) : nothing
Comment thread
kshyatt marked this conversation as resolved.
dret = if !isa(t, Const) && EnzymeRules.needs_shadow(config)
block(t.dval, c.val)
elseif EnzymeRules.needs_shadow(config)
Enzyme.make_zero(ret)
else
nothing
end
return EnzymeRules.AugmentedReturn(ret, dret, dret)
end

function EnzymeRules.reverse(
config::EnzymeRules.RevConfigWidth{1},
::Const{typeof(block)},
::Type{RT},
cache,
t::Annotation{<:AbstractTensorMap},
c::Annotation{<:Sector},
) where {RT}
Comment thread
lkdvos marked this conversation as resolved.
return (nothing, nothing)
end

Expand Down Expand Up @@ -78,9 +103,38 @@ function EnzymeRules.forward(
end
end

function EnzymeRules.forward(
config::EnzymeRules.FwdConfigWidth{1},
::Const{typeof(block)},
::Type{RT},
t::Annotation{<:AbstractTensorMap},
c::Annotation{<:Sector},
) where {RT}
ret = EnzymeRules.needs_primal(config) ? block(t.val, c.val) : nothing
dret = if !isa(t, Const) && EnzymeRules.needs_shadow(config)
block(t.dval, c.val)
elseif EnzymeRules.needs_shadow(config)
Enzyme.make_zero(ret)
else
nothing
end
if EnzymeRules.needs_primal(config) && EnzymeRules.needs_shadow(config)
return Duplicated(ret, dret)
elseif EnzymeRules.needs_primal(config)
return ret
elseif EnzymeRules.needs_shadow(config)
return dret
else
return nothing
end
end

@inline EnzymeRules.inactive(::typeof(TensorKit.fsbraid), ::Any) = nothing
@inline EnzymeRules.inactive(::typeof(TensorKit.fsbraid), ::Any, ::Any) = nothing
@inline EnzymeRules.inactive(::typeof(TensorKit.artin_braid), ::Any, ::Any) = nothing
@inline EnzymeRules.inactive(::typeof(TensorKit.insertleftunit), ::HomSpace, ::Any) = nothing
@inline EnzymeRules.inactive(::typeof(TensorKit.insertrightunit), ::HomSpace, ::Any) = nothing
@inline EnzymeRules.inactive(::typeof(TensorKit.removeunit), ::HomSpace, ::Any) = nothing
@inline EnzymeRules.inactive(::typeof(TensorKit.sectorstructure), ::Any) = nothing
@inline EnzymeRules.inactive(::typeof(TensorKit.degeneracystructure), ::Any) = nothing
@inline EnzymeRules.inactive(::typeof(TensorKit.select), s::HomSpace, i::Index2Tuple) = nothing
Expand Down
30 changes: 21 additions & 9 deletions test/enzyme-indexmanipulations-unit/insertunit.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,21 +6,33 @@ using Random
spacelist = ad_spacelist(fast_tests)
eltypes = (Float64, ComplexF64)

if VERSION > v"1.11.0-rc" # https://github.com/QuantumKitHub/TensorKit.jl/issues/457
@timedtestset verbose = true "Enzyme - Index Manipulations (insertunit):" begin
@timedtestset verbose = true "$(TensorKit.type_repr(sectortype(eltype(V)))) ($T)" for V in spacelist, T in eltypes, TA in (Duplicated,)
atol = default_tol(T)
rtol = default_tol(T)
A = randn(T, V[1] ⊗ V[2] ← (V[3] ⊗ V[4] ⊗ V[5])')
@testset for insertunit in (insertleftunit, insertrightunit)
EnzymeTestUtils.test_reverse(insertunit, TA, (A, TA), (Val(1), Const); atol, rtol)
EnzymeTestUtils.test_reverse(insertunit, TA, (A, TA), (Val(4), Const); atol, rtol)
is_ci = get(ENV, "CI", "false") == "true"

@timedtestset verbose = true "Enzyme - Index Manipulations (insertunit):" begin
@timedtestset verbose = true "$(TensorKit.type_repr(sectortype(eltype(V)))) ($T)" for V in spacelist, T in eltypes, TA in (Duplicated,)
atol = default_tol(T)
rtol = default_tol(T)
A = randn(T, V[1] ⊗ V[2] ← (V[3] ⊗ V[4] ⊗ V[5])')
@testset for insertunit in (insertleftunit, insertrightunit)
EnzymeTestUtils.test_reverse(insertunit, TA, (A, TA), (Val(1), Const); atol, rtol)
EnzymeTestUtils.test_reverse(insertunit, TA, (A, TA), (Val(4), Const); atol, rtol)
if !is_ci
EnzymeTestUtils.test_reverse(insertunit, TA, (A', TA), (Val(2), Const); atol, rtol)
EnzymeTestUtils.test_reverse(insertunit, TA, (A, TA), (Val(1), Const); atol, rtol, fkwargs = (copy = false,))
EnzymeTestUtils.test_reverse(insertunit, TA, (A, TA), (Val(2), Const); atol, rtol, fkwargs = (copy = true,))
EnzymeTestUtils.test_reverse(insertunit, TA, (A, TA), (Val(3), Const); atol, rtol, fkwargs = (copy = false, dual = true, conj = true))
EnzymeTestUtils.test_reverse(insertunit, TA, (A', TA), (Val(3), Const); atol, rtol, fkwargs = (copy = false, dual = true, conj = true))
end

EnzymeTestUtils.test_forward(insertunit, TA, (A, TA), (Val(1), Const); atol, rtol)
EnzymeTestUtils.test_forward(insertunit, TA, (A, TA), (Val(4), Const); atol, rtol)
if !is_ci
EnzymeTestUtils.test_forward(insertunit, TA, (A', TA), (Val(2), Const); atol, rtol)
EnzymeTestUtils.test_forward(insertunit, TA, (A, TA), (Val(1), Const); atol, rtol, fkwargs = (copy = false,))
EnzymeTestUtils.test_forward(insertunit, TA, (A, TA), (Val(2), Const); atol, rtol, fkwargs = (copy = true,))
EnzymeTestUtils.test_forward(insertunit, TA, (A, TA), (Val(3), Const); atol, rtol, fkwargs = (copy = false, dual = true, conj = true))
EnzymeTestUtils.test_forward(insertunit, TA, (A', TA), (Val(3), Const); atol, rtol, fkwargs = (copy = false, dual = true, conj = true))
end
end
end
end
23 changes: 12 additions & 11 deletions test/enzyme-indexmanipulations-unit/removeunit.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,17 +6,18 @@ using Random
spacelist = ad_spacelist(fast_tests)
eltypes = (Float64, ComplexF64)

if VERSION > v"1.11.0-rc" # https://github.com/QuantumKitHub/TensorKit.jl/issues/457
@timedtestset verbose = true "Enzyme - Index Manipulations (removeunit):" begin
@timedtestset verbose = true "$(TensorKit.type_repr(sectortype(eltype(V)))) ($T)" for V in spacelist, T in eltypes, TB in (Duplicated,)
atol = default_tol(T)
rtol = default_tol(T)
A = randn(T, V[1] ⊗ V[2] ← (V[3] ⊗ V[4] ⊗ V[5])')
for i in 1:2
B = insertleftunit(A, i; dual = rand(Bool))
EnzymeTestUtils.test_reverse(removeunit, TB, (B, TB), (Val(i), Const); atol, rtol, fkwargs = (copy = false,))
EnzymeTestUtils.test_reverse(removeunit, TB, (B, TB), (Val(i), Const); atol, rtol, fkwargs = (copy = true,))
end
@timedtestset verbose = true "Enzyme - Index Manipulations (removeunit):" begin
@timedtestset verbose = true "$(TensorKit.type_repr(sectortype(eltype(V)))) ($T)" for V in spacelist, T in eltypes, TB in (Duplicated,)
atol = default_tol(T)
rtol = default_tol(T)
A = randn(T, V[1] ⊗ V[2] ← (V[3] ⊗ V[4] ⊗ V[5])')
for i in 1:2
B = insertleftunit(A, i; dual = rand(Bool))
EnzymeTestUtils.test_reverse(removeunit, TB, (B, TB), (Val(i), Const); atol, rtol, fkwargs = (copy = false,))
EnzymeTestUtils.test_reverse(removeunit, TB, (B, TB), (Val(i), Const); atol, rtol, fkwargs = (copy = true,))

EnzymeTestUtils.test_forward(removeunit, TB, (B, TB), (Val(i), Const); atol, rtol, fkwargs = (copy = false,))
EnzymeTestUtils.test_forward(removeunit, TB, (B, TB), (Val(i), Const); atol, rtol, fkwargs = (copy = true,))
end
end
end
Loading