diff --git a/ext/TensorKitEnzymeExt/indexmanipulations.jl b/ext/TensorKitEnzymeExt/indexmanipulations.jl index 20439adaa..3a54a3230 100644 --- a/ext/TensorKitEnzymeExt/indexmanipulations.jl +++ b/ext/TensorKitEnzymeExt/indexmanipulations.jl @@ -183,57 +183,6 @@ function EnzymeRules.forward( end end -# Differentiating through the fusion tree loop corrupts the -# heap on Julia 1.10 and causes segfaults in the GC. Remove this -# custom rule when we drop support for 1.10 -function EnzymeRules.augmented_primal( - config::EnzymeRules.RevConfigWidth{1}, - func::Const{typeof(twist!)}, - ::Type{RT}, - t::Annotation{<:AbstractTensorMap}, - inds::Const; - inv::Bool = false - ) where {RT} - twist!(t.val, inds.val; inv) - primal = EnzymeRules.needs_primal(config) ? t.val : nothing - shadow = EnzymeRules.needs_shadow(config) ? t.dval : nothing - return EnzymeRules.AugmentedReturn(primal, shadow, nothing) -end - -function EnzymeRules.reverse( - config::EnzymeRules.RevConfigWidth{1}, - func::Const{typeof(twist!)}, - ::Type{RT}, - cache, - t::Annotation{<:AbstractTensorMap}, - inds::Const; - inv::Bool = false - ) where {RT} - !isa(t, Const) && twist!(t.dval, inds.val; inv = !inv) - return (nothing, nothing) -end - -function EnzymeRules.forward( - config::EnzymeRules.FwdConfigWidth{1}, - func::Const{typeof(twist!)}, - ::Type{RT}, - t::Annotation{<:AbstractTensorMap}, - inds::Annotation; - inv::Bool = false - ) where {RT} - twist!(t.val, inds.val; inv) - !isa(t, Const) && twist!(t.dval, inds.val; inv) - if EnzymeRules.needs_primal(config) && EnzymeRules.needs_shadow(config) - return Duplicated(t.val, t.dval) - elseif EnzymeRules.needs_primal(config) - return t.val - elseif EnzymeRules.needs_shadow(config) - return t.dval - else - return nothing - end -end - function EnzymeRules.augmented_primal( config::EnzymeRules.RevConfigWidth{1}, func::Const{typeof(flip)}, diff --git a/ext/TensorKitEnzymeExt/linalg.jl b/ext/TensorKitEnzymeExt/linalg.jl index 232e58303..8401575e7 100644 --- a/ext/TensorKitEnzymeExt/linalg.jl +++ b/ext/TensorKitEnzymeExt/linalg.jl @@ -125,30 +125,6 @@ function EnzymeRules.reverse( ) return (nothing,) end -function EnzymeRules.forward( - config::EnzymeRules.FwdConfigWidth{1}, - ::Type{RT}, - func::Const{typeof(tr)}, - A::Annotation{<:AbstractTensorMap}, - ) where {RT} - y = EnzymeRules.needs_primal(config) ? tr(A.val) : nothing - Δy = if EnzymeRules.needs_shadow(config) && !isa(A, Const) - tr(A.dval) - elseif EnzymeRules.needs_shadow(config) - zero(eltype(A.dval)) - else - nothing - end - if EnzymeRules.needs_primal(config) && EnzymeRules.needs_shadow(config) - return Duplicated(y, Δy) - elseif EnzymeRules.needs_primal(config) - return y - elseif EnzymeRules.needs_shadow(config) - return Δy - else - return nothing - end -end function EnzymeRules.augmented_primal( config::EnzymeRules.RevConfigWidth{1}, func::Const{typeof(norm)},