From 0e8d4b03b762e5aec1a41816f3288eaabf4d3d38 Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Mon, 13 Jul 2026 11:23:06 -0400 Subject: [PATCH 1/2] Does Enzyme need a custom rule for trace_permute --- ext/TensorKitEnzymeExt/tensoroperations.jl | 89 ---------------------- 1 file changed, 89 deletions(-) diff --git a/ext/TensorKitEnzymeExt/tensoroperations.jl b/ext/TensorKitEnzymeExt/tensoroperations.jl index e7170b6c5..c522081a3 100644 --- a/ext/TensorKitEnzymeExt/tensoroperations.jl +++ b/ext/TensorKitEnzymeExt/tensoroperations.jl @@ -117,92 +117,3 @@ function EnzymeRules.forward( end end -# tensortrace! -# ------------ - -function EnzymeRules.augmented_primal( - config::EnzymeRules.RevConfigWidth{1}, - func::Const{typeof(TensorKit.trace_permute!)}, - ::Type{RT}, - C::Annotation{<:AbstractTensorMap}, - A::Annotation{<:AbstractTensorMap}, - p::Const{<:Index2Tuple}, - q::Const{<:Index2Tuple}, - α::Annotation{<:Number}, - β::Annotation{<:Number}, - backend::Const, - ) where {RT} - C_cache = !isa(β, Const) ? copy(C.val) : nothing - A_cache = EnzymeRules.overwritten(config)[3] ? copy(A.val) : nothing - At = if !isa(α, Const) - At = TO.tensortrace(A.val, p.val, q.val, false, One(), backend.val) - add!(C.val, At, α.val, β.val) - At - else - TensorKit.trace_permute!(C.val, A.val, p.val, q.val, α.val, β.val, backend.val) - nothing - end - primal = EnzymeRules.needs_primal(config) ? C.val : nothing - shadow = EnzymeRules.needs_shadow(config) ? C.dval : nothing - cache = (C_cache, A_cache, At) - return EnzymeRules.AugmentedReturn(primal, shadow, cache) -end - - -function EnzymeRules.reverse( - config::EnzymeRules.RevConfigWidth{1}, - func::Const{typeof(TensorKit.trace_permute!)}, - ::Type{RT}, - cache, - C::Annotation{<:AbstractTensorMap}, - A::Annotation{<:AbstractTensorMap}, - p::Const{<:Index2Tuple}, - q::Const{<:Index2Tuple}, - α::Annotation{<:Number}, - β::Annotation{<:Number}, - backend::Const, - ) where {RT} - C_cache, A_cache, At = cache - Aval = something(A_cache, A.val) - Cval = something(C_cache, C.val) - !isa(A, Const) && !isa(C, Const) && TensorKit.trace_permute_pullback_ΔA!(A.dval, C.dval, Aval, p.val, q.val, α.val, backend.val) - Δαr = pullback_dα(α, C, At) - Δβr = pullback_dβ(β, C, Cval) - !isa(C, Const) && pullback_dC!(C.dval, β.val) - return nothing, nothing, nothing, nothing, Δαr, Δβr, nothing -end - -function EnzymeRules.forward( - config::EnzymeRules.FwdConfigWidth{1}, - func::Const{typeof(TensorKit.trace_permute!)}, - ::Type{RT}, - C::Annotation{<:AbstractTensorMap}, - A::Annotation{<:AbstractTensorMap}, - p::Annotation{<:Index2Tuple}, - q::Annotation{<:Index2Tuple}, - α::Annotation{<:Number}, - β::Annotation{<:Number}, - backend::Const, - ) where {RT} - # dD = dα * tr(A) + α * tr(dA) + dβ * C + β * dC - # dC1 = dβ * C + β * dC - if !isa(C, Const) - if isa(β, Const) - scale!(C.dval, β.val) - else - add!(C.dval, C.val, β.dval, β.val) - end - !isa(α, Const) && TensorKit.trace_permute!(C.dval, A.val, p.val, q.val, α.dval, One(), backend.val) - !isa(A, Const) && TensorKit.trace_permute!(C.dval, A.dval, p.val, q.val, α.val, One(), backend.val) - end - TensorKit.trace_permute!(C.val, A.val, p.val, q.val, α.val, β.val, backend.val) - if EnzymeRules.needs_primal(config) && EnzymeRules.needs_shadow(config) - return C - elseif EnzymeRules.needs_primal(config) - return C.val - elseif EnzymeRules.needs_shadow(config) - return C.dval - else - return nothing - end -end From 61eba8a323b8ca89b737e4d9cac4592c8b22eb5d Mon Sep 17 00:00:00 2001 From: Katharine Hyatt Date: Tue, 14 Jul 2026 14:18:34 +0200 Subject: [PATCH 2/2] Working --- ext/TensorKitEnzymeExt/TensorKitEnzymeExt.jl | 1 + ext/TensorKitEnzymeExt/tensoroperations.jl | 1 - ext/TensorKitEnzymeExt/utility.jl | 62 ++++++++++++++++++++ 3 files changed, 63 insertions(+), 1 deletion(-) diff --git a/ext/TensorKitEnzymeExt/TensorKitEnzymeExt.jl b/ext/TensorKitEnzymeExt/TensorKitEnzymeExt.jl index ba628b67f..6b3e9c283 100644 --- a/ext/TensorKitEnzymeExt/TensorKitEnzymeExt.jl +++ b/ext/TensorKitEnzymeExt/TensorKitEnzymeExt.jl @@ -3,6 +3,7 @@ module TensorKitEnzymeExt using Enzyme using TensorKit import TensorKit as TK +using TensorKit: subblock using VectorInterface using TensorOperations: TensorOperations, IndexTuple, Index2Tuple, linearize import TensorOperations as TO diff --git a/ext/TensorKitEnzymeExt/tensoroperations.jl b/ext/TensorKitEnzymeExt/tensoroperations.jl index c522081a3..8d6652ba7 100644 --- a/ext/TensorKitEnzymeExt/tensoroperations.jl +++ b/ext/TensorKitEnzymeExt/tensoroperations.jl @@ -116,4 +116,3 @@ function EnzymeRules.forward( return nothing end end - diff --git a/ext/TensorKitEnzymeExt/utility.jl b/ext/TensorKitEnzymeExt/utility.jl index 2d6366c90..a364c5daf 100644 --- a/ext/TensorKitEnzymeExt/utility.jl +++ b/ext/TensorKitEnzymeExt/utility.jl @@ -19,6 +19,68 @@ pullback_dC!(ΔC, β::Number) = scale!(ΔC, conj(β)) @inline EnzymeRules.inactive_type(::Type{<:TensorKit.GenericTreeTransformer}) = true @inline EnzymeRules.inactive_type(::Type{<:TensorKit.VectorSpace}) = true +function EnzymeRules.augmented_primal( + config::EnzymeRules.RevConfigWidth{1}, + ::Const{typeof(subblock)}, + ::Type{RT}, + t::Annotation{<:AbstractTensorMap}, + f::Annotation, + ) where {RT} + ret = EnzymeRules.needs_primal(config) ? subblock(t.val, f.val) : nothing + dret = if !isa(t, Const) && EnzymeRules.needs_shadow(config) + subblock(t.dval, f.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(subblock)}, + ::Type{RT}, + cache, + 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 + +function EnzymeRules.forward( + config::EnzymeRules.FwdConfigWidth{1}, + ::Const{typeof(subblock)}, + ::Type{RT}, + t::Annotation{<:AbstractTensorMap}, + f::Annotation, + ) where {RT} + ret = EnzymeRules.needs_primal(config) ? subblock(t.val, f.val) : nothing + dret = if !isa(t, Const) && EnzymeRules.needs_shadow(config) + subblock(t.dval, f.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.sectorstructure), ::Any) = nothing @inline EnzymeRules.inactive(::typeof(TensorKit.degeneracystructure), ::Any) = nothing @inline EnzymeRules.inactive(::typeof(TensorKit.select), s::HomSpace, i::Index2Tuple) = nothing