From c157be09e8e348d8a289a68d4d1c7af6db203014 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Wed, 15 Jul 2026 07:11:14 +0200 Subject: [PATCH 1/3] Try letting Enzyme make rules for insert/remove unit --- ext/TensorKitEnzymeExt/TensorKitEnzymeExt.jl | 2 +- ext/TensorKitEnzymeExt/indexmanipulations.jl | 89 ------------------- ext/TensorKitEnzymeExt/utility.jl | 62 +++++++++++++ .../insertunit.jl | 36 ++++---- .../removeunit.jl | 23 ++--- 5 files changed, 96 insertions(+), 116 deletions(-) diff --git a/ext/TensorKitEnzymeExt/TensorKitEnzymeExt.jl b/ext/TensorKitEnzymeExt/TensorKitEnzymeExt.jl index 6b3e9c283..cadae496a 100644 --- a/ext/TensorKitEnzymeExt/TensorKitEnzymeExt.jl +++ b/ext/TensorKitEnzymeExt/TensorKitEnzymeExt.jl @@ -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 diff --git a/ext/TensorKitEnzymeExt/indexmanipulations.jl b/ext/TensorKitEnzymeExt/indexmanipulations.jl index f772348e5..0adbb0600 100644 --- a/ext/TensorKitEnzymeExt/indexmanipulations.jl +++ b/ext/TensorKitEnzymeExt/indexmanipulations.jl @@ -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 diff --git a/ext/TensorKitEnzymeExt/utility.jl b/ext/TensorKitEnzymeExt/utility.jl index a364c5daf..618cf59cc 100644 --- a/ext/TensorKitEnzymeExt/utility.jl +++ b/ext/TensorKitEnzymeExt/utility.jl @@ -52,6 +52,39 @@ function EnzymeRules.reverse( 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 + 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} + dret = cache + if !isnothing(dret) && !isa(t, Const) + block(t.dval, c.val) .= dret + end + return (nothing, nothing) +end + function EnzymeRules.forward( config::EnzymeRules.FwdConfigWidth{1}, ::Const{typeof(subblock)}, @@ -78,9 +111,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 diff --git a/test/enzyme-indexmanipulations-unit/insertunit.jl b/test/enzyme-indexmanipulations-unit/insertunit.jl index 9513d5ee7..6e89655d1 100644 --- a/test/enzyme-indexmanipulations-unit/insertunit.jl +++ b/test/enzyme-indexmanipulations-unit/insertunit.jl @@ -6,21 +6,27 @@ 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) - 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 +@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) + 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)) + + EnzymeTestUtils.test_forward(insertunit, TA, (A, TA), (Val(1), Const); atol, rtol) + EnzymeTestUtils.test_forward(insertunit, TA, (A, TA), (Val(4), Const); atol, rtol) + 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 diff --git a/test/enzyme-indexmanipulations-unit/removeunit.jl b/test/enzyme-indexmanipulations-unit/removeunit.jl index 2cd632a90..a5b42d0bb 100644 --- a/test/enzyme-indexmanipulations-unit/removeunit.jl +++ b/test/enzyme-indexmanipulations-unit/removeunit.jl @@ -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 From a411580a2936f2ee8aa35eb30b967406975b975f Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Wed, 15 Jul 2026 10:36:30 +0200 Subject: [PATCH 2/3] Run fewer insert tests --- .../insertunit.jl | 26 ++++++++++++------- 1 file changed, 16 insertions(+), 10 deletions(-) diff --git a/test/enzyme-indexmanipulations-unit/insertunit.jl b/test/enzyme-indexmanipulations-unit/insertunit.jl index 6e89655d1..bd59f9d62 100644 --- a/test/enzyme-indexmanipulations-unit/insertunit.jl +++ b/test/enzyme-indexmanipulations-unit/insertunit.jl @@ -6,6 +6,8 @@ using Random spacelist = ad_spacelist(fast_tests) eltypes = (Float64, ComplexF64) +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) @@ -14,19 +16,23 @@ eltypes = (Float64, ComplexF64) @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) - 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)) + 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) - 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)) + 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 From 64010411875fbcb5bbd5de129c1200993ae94678 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Wed, 15 Jul 2026 11:15:59 +0200 Subject: [PATCH 3/3] Get rid of unneeded copies --- ext/TensorKitEnzymeExt/utility.jl | 8 -------- 1 file changed, 8 deletions(-) diff --git a/ext/TensorKitEnzymeExt/utility.jl b/ext/TensorKitEnzymeExt/utility.jl index 618cf59cc..dd99d96e8 100644 --- a/ext/TensorKitEnzymeExt/utility.jl +++ b/ext/TensorKitEnzymeExt/utility.jl @@ -45,10 +45,6 @@ function EnzymeRules.reverse( t::Annotation{<:AbstractTensorMap}, f::Annotation, ) where {RT} - dret = cache - if !isnothing(dret) && !isa(t, Const) - subblock(t.dval, f.val) .= dret - end return (nothing, nothing) end @@ -78,10 +74,6 @@ function EnzymeRules.reverse( t::Annotation{<:AbstractTensorMap}, c::Annotation{<:Sector}, ) where {RT} - dret = cache - if !isnothing(dret) && !isa(t, Const) - block(t.dval, c.val) .= dret - end return (nothing, nothing) end