proper fix (hopefullly)

Author: Jutho

ext/TensorKitChainRulesCoreExt/tensoroperations.jl

@@ -61,7 +61,8 @@ function ChainRulesCore.rrule(::typeof(TensorOperations.tensorcontract!),
         pΔC = _repartition(ipAB, TO.numout(pA))
 
         dC = @thunk projectC(scale(ΔC, conj(β)))
-        dA = @thunk let ipA = _repartition(invperm(linearize(pA)), A)
+        dA = @thunk let
+            ipA = _repartition(invperm(linearize(pA)), A)
             conjΔC = conjA
             conjB′ = conjA ? conjB : !conjB
             TA = promote_contract(scalartype(ΔC), scalartype(B), scalartype(α))
@@ -77,7 +78,8 @@ function ChainRulesCore.rrule(::typeof(TensorOperations.tensorcontract!),
                                   conjA ? α : conj(α), Zero(), ba...)
             return projectA(_dA)
         end
-        dB = @thunk let ipB = _repartition(invperm(linearize(pB)), B)
+        dB = @thunk let
+            ipB = _repartition(invperm(linearize(pB)), B)
             conjΔC = conjB
             conjA′ = conjB ? conjA : !conjA
             TB = promote_contract(scalartype(ΔC), scalartype(A), scalartype(α))
@@ -123,7 +125,9 @@ function ChainRulesCore.rrule(::typeof(TensorOperations.tensortrace!),
     function pullback(ΔC′)
         ΔC = unthunk(ΔC′)
         dC = @thunk projectC(scale(ΔC, conj(β)))
-        dA = @thunk let ip = invperm((linearize(p)..., q[1]..., q[2]...)), pdA = _repartition(ip, A)
+        dA = @thunk let
+            ip = invperm((linearize(p)..., q[1]..., q[2]...))
+            pdA = _repartition(ip, A)
             E = one!(TO.tensoralloc_add(scalartype(A), A, q, conjA))
             twist!(E, filter(x -> !isdual(space(E, x)), codomainind(E)))
             pE = ((), trivtuple(TO.numind(q)))

ext/TensorKitChainRulesCoreExt/utility.jl

@@ -2,13 +2,15 @@
 # -------
 trivtuple(N) = ntuple(identity, N)
 
-function _repartition(p::IndexTuple, N₁::Int)
+Base.@constprop :aggressive function _repartition(p::IndexTuple, N₁::Int)
     length(p) >= N₁ ||
         throw(ArgumentError("cannot repartition $(typeof(p)) to $N₁, $(length(p) - N₁)"))
     return TupleTools.getindices(p, trivtuple(N₁)),
            TupleTools.getindices(p, trivtuple(length(p) - N₁) .+ N₁)
 end
-_repartition(p::Index2Tuple, N₁::Int) = _repartition(linearize(p), N₁)
+Base.@constprop :aggressive function _repartition(p::Index2Tuple, N₁::Int)
+    return _repartition(linearize(p), N₁)
+end
 function _repartition(p::Union{IndexTuple,Index2Tuple}, ::Index2Tuple{N₁}) where {N₁}
     return _repartition(p, N₁)
 end