Enzyme doesn't need a custom rule for trace_permute

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

ext/TensorKitEnzymeExt/tensoroperations.jl

@@ -116,93 +116,3 @@ function EnzymeRules.forward(
         return nothing
     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

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