Does Enzyme need a custom rule for trace_permute

Author: kshyatt

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