From 2b08efee612f8bc51629454b940480cc60311a63 Mon Sep 17 00:00:00 2001
From: lkdvos <ldevos98@gmail.com>
Date: Mon, 12 Oct 2026 16:01:08 -0400
Subject: [PATCH 1/7] Specialize `makeblascontractable` on the scalar type

`makeblascontractable` received the scalar type of the destination as an unconstrained
argument, which Julia does not specialize on. As a result, the allocation and permutation of
inputs that are not directly BLAS-contractable were dynamically dispatched, which also
affects the `@tensor` macro.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 docs/src/changelog.md              | 1 +
 src/implementation/blascontract.jl | 2 +-
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/docs/src/changelog.md b/docs/src/changelog.md
index c8f910b3..6bd85d79 100644
--- a/docs/src/changelog.md
+++ b/docs/src/changelog.md
@@ -19,6 +19,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Performance
 
+- Contractions of `Array`s that require a permuted copy of an input no longer suffer from a type instability.
 - `ncon` partitions its intermediates as `(open indices of A; open indices of B)` instead of placing all indices in the codomain, which avoids repartitioning copies for tensor types such as `TensorMap`s with symmetries ([#311](https://github.com/QuantumKitHub/TensorOperations.jl/issues/311)).
 
 ## [5.8.2](https://github.com/QuantumKitHub/TensorOperations.jl/compare/v5.8.1...v5.8.2) - 2026-09-30
diff --git a/src/implementation/blascontract.jl b/src/implementation/blascontract.jl
index 72bb9278..289ae0c2 100644
--- a/src/implementation/blascontract.jl
+++ b/src/implementation/blascontract.jl
@@ -97,7 +97,7 @@ function _unsafe_blas_contract!(
     return C
 end
 
-function makeblascontractable(A, pA, TC, backend, allocator)
+function makeblascontractable(A, pA, ::Type{TC}, backend, allocator) where {TC}
     flagA = isblascontractable(A, pA) && eltype(A) == TC
     if !flagA
         A_ = tensoralloc_add(TC, A, pA, false, Val(true), allocator)

From 5b340cb279dcd43fa8d58871c3ab6552606af0bb Mon Sep 17 00:00:00 2001
From: lkdvos <ldevos98@gmail.com>
Date: Mon, 12 Oct 2026 16:02:53 -0400
Subject: [PATCH 2/7] Specialize the functions that analyze `ncon` networks

The index notation files are included in a module-level `@nospecialize` block to reduce the
compilation time of the macros. However, `ncontree.jl` and `indexordertree.jl` contain
functions that are also called at runtime by `ncon`, such as `check_nconstyle`, `ncontree`
and `indexordertree`, which therefore ran without specialization on their arguments.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/TensorOperations.jl | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/TensorOperations.jl b/src/TensorOperations.jl
index 499dbe30..375d6bb3 100644
--- a/src/TensorOperations.jl
+++ b/src/TensorOperations.jl
@@ -51,8 +51,11 @@ include("indexnotation/analyzers.jl")
 include("indexnotation/preprocessors.jl")
 include("indexnotation/postprocessors.jl")
 include("indexnotation/contractiontrees.jl")
+# also used at runtime by `ncon`, so these should be specialized
+@specialize
 include("indexnotation/ncontree.jl")
 include("indexnotation/indexordertree.jl")
+@nospecialize
 include("indexnotation/poly.jl")
 include("indexnotation/optdata.jl")
 include("indexnotation/optimaltree.jl")

From 66682154e18b1d3aaea53b351faecafd17604cf6 Mon Sep 17 00:00:00 2001
From: lkdvos <ldevos98@gmail.com>
Date: Mon, 12 Oct 2026 16:04:23 -0400
Subject: [PATCH 3/7] Faster validation and contraction tree construction for
 `ncon` networks

- `check_nconstyle` and `_verify_indexorder` count label occurrences instead of repeatedly
  searching and deleting from an abstractly typed vector.
- `ncontree` and `indexordertree` share a single implementation that merges partial trees
  in a type-stable manner, and represent the contraction tree as a list of pairwise
  contractions. The nested trees that are returned are unchanged.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 src/indexnotation/indexordertree.jl |  72 +++++-----------
 src/indexnotation/ncontree.jl       | 122 +++++++++++++++++++---------
 2 files changed, 104 insertions(+), 90 deletions(-)

diff --git a/src/indexnotation/indexordertree.jl b/src/indexnotation/indexordertree.jl
index c1ad07bc..2a87a6af 100644
--- a/src/indexnotation/indexordertree.jl
+++ b/src/indexnotation/indexordertree.jl
@@ -1,63 +1,33 @@
 # verify that every conraction index appears exactly twice, and any other index exactly once
 function _verify_indexorder(network, indexorder)
-    allindices = Vector{Any}()
-    for ind in network
-        append!(allindices, ind)
-    end
-    while length(allindices) > 0
-        i = pop!(allindices)
-        if i ∈ indexorder # contraction index
-            k = findfirst(isequal(i), allindices)
-            k === nothing && return false
-            l = findnext(isequal(i), allindices, k + 1)
-            l !== nothing && return false
-            deleteat!(allindices, k)
-        else
-            findfirst(isequal(i), allindices) === nothing || return false
+    for indices in network
+        for i in indices
+            n = sum(ind -> count(isequal(i), ind), network)
+            n == (i in indexorder ? 2 : 1) || return false
         end
     end
     return true
 end
 
 function indexordertree(network, indexorder)
-    _verify_indexorder(network, indexorder) ||
-        throw(ArgumentError("not a valid contraction tree with contraction indices $indexorder"))
-    contractionindices = Vector{Vector{Any}}(undef, length(network))
-    for k in 1:length(network)
-        indices = network[k]
-        # trace indices have already been removed, remove open indices by filtering on positive values
-        contractionindices[k] = filter(in(indexorder), indices)
-    end
-    partialtrees = collect(Any, 1:length(network))
-    return _indexordertree!(partialtrees, contractionindices, indexorder)
+    return _nestedtree(_indexordertree_steps(network, indexorder), length(network))
 end
 
-function _indexordertree!(partialtrees, contractionindices, indexorder)
-    while length(indexorder) > 0
-        c = indexorder[1]
-        i1 = findfirst(x -> in(c, x), contractionindices)
-        if i1 === nothing
-            indexorder = indexorder[2:end]
-        else
-            i2 = findnext(x -> in(c, x), contractionindices, i1 + 1)
-            @assert i2 !== nothing
-            cs = intersect(contractionindices[i1], contractionindices[i2])
-            indexorder = setdiff(indexorder, cs)
-            newindices = unique2(vcat(contractionindices[i1], contractionindices[i2]))
-            newtree = Any[partialtrees[i1], partialtrees[i2]]
-            partialtrees[i1] = newtree
-            deleteat!(partialtrees, i2)
-            contractionindices[i1] = newindices
-            deleteat!(contractionindices, i2)
-        end
-    end
-    if length(partialtrees) > 1
-        @assert all(isempty, contractionindices) # disconnected network
-        while length(partialtrees) > 1
-            partialtrees[end - 1] = Any[partialtrees[end - 1], partialtrees[end]]
-            pop!(partialtrees)
-            pop!(contractionindices)
-        end
+# Determine the contraction order where the indices are contracted in the order specified by
+# `indexorder`, as a list of pairwise contractions (see `_mergetrees!`).
+function _indexordertree_steps(network, indexorder)
+    _verify_indexorder(network, indexorder) ||
+        throw(ArgumentError("not a valid contraction tree with contraction indices $indexorder"))
+    steps = NTuple{2, Int}[]
+    partialtrees = collect(1:length(network))
+    group = collect(1:length(network))
+    for c in indexorder
+        k1 = findfirst(indices -> c in indices, network)
+        isnothing(k1) && continue
+        # trace indices have already been removed, so `c` appears in another tensor
+        k2 = findnext(indices -> c in indices, network, k1 + 1)
+        isnothing(k2) && continue
+        _mergetrees!(steps, partialtrees, group, k1, k2)
     end
-    return partialtrees[1]
+    return _finalizetree!(steps, partialtrees, group)
 end
diff --git a/src/indexnotation/ncontree.jl b/src/indexnotation/ncontree.jl
index ed8f6f88..b16cee58 100644
--- a/src/indexnotation/ncontree.jl
+++ b/src/indexnotation/ncontree.jl
@@ -5,63 +5,107 @@ check_nconstyle(network) = _check_nconstyle(network, Val(false))
 function _check_nconstyle(network, ::Val{check}) where {check}
     allindices = Vector{Int}()
     for ind in network
-        all(i -> isa(i, Integer), ind) || return check ? false :
-            throw(IndexError("All indices must be integers"))
-        append!(allindices, ind)
+        for i in ind
+            isa(i, Integer) || return check ? false :
+                throw(IndexError("All indices must be integers"))
+            push!(allindices, i)
+        end
     end
-    while length(allindices) > 0
-        i = pop!(allindices)
+    sort!(allindices)
+    j = 1
+    while j <= length(allindices)
+        i = allindices[j]
+        # number of times that `i` appears
+        n = 1
+        while j + n <= length(allindices) && allindices[j + n] == i
+            n += 1
+        end
         if i > 0 # positive labels represent contractions or traces and should appear twice
-            k = findfirst(isequal(i), allindices)
-            isnothing(k) && return check ? false :
+            n == 1 && return check ? false :
                 throw(IndexError(lazy"Index $i appears only once in the network"))
-            l = findnext(isequal(i), allindices, k + 1)
-            !isnothing(l) && return check ? false :
+            n > 2 && return check ? false :
                 throw(IndexError(lazy"Index $i appears more than twice in the network"))
-            deleteat!(allindices, k)
         elseif i < 0 # negative labels represent open indices and should appear once
-            isnothing(findfirst(isequal(i), allindices)) || return check ? false :
+            n > 1 && return check ? false :
                 throw(IndexError(lazy"Index $i appears more than once in the network"))
         else # i == 0
             return check ? false : throw(IndexError("Index 0 is not allowed in the network"))
         end
+        j += n
     end
     return check ? true : nothing
 end
 
-function ncontree(network)
-    contractionindices = Vector{Vector{Int}}(undef, length(network))
-    for k in 1:length(network)
-        indices = network[k]
-        # trace indices have already been removed, remove open indices by filtering on positive values
-        contractionindices[k] = filter(i -> i > 0, indices)
+ncontree(network) = _nestedtree(_ncontree_steps(network), length(network))
+
+# Determine the contraction order where the indices are contracted in increasing order, as a
+# list of pairwise contractions (see `_mergetrees!`).
+function _ncontree_steps(network)
+    # collect the positive labels, i.e. the contraction indices, and the tensors they belong to
+    contractionindices = Int[]
+    owners = Int[]
+    for (k, indices) in enumerate(network)
+        for i in indices
+            i > 0 || continue
+            push!(contractionindices, i)
+            push!(owners, k)
+        end
     end
-    partialtrees = collect(Any, 1:length(network))
-    return _ncontree!(partialtrees, contractionindices)
+
+    # contract the indices in increasing order
+    steps = NTuple{2, Int}[]
+    partialtrees = collect(1:length(network))
+    group = collect(1:length(network))
+    for j in sortperm(contractionindices)
+        k = findnext(==(contractionindices[j]), contractionindices, j + 1)
+        # each index appears twice, only consider the first occurrence
+        isnothing(k) || _mergetrees!(steps, partialtrees, group, owners[j], owners[k])
+    end
+    return _finalizetree!(steps, partialtrees, group)
 end
 
-function _ncontree!(partialtrees, contractionindices)
-    if length(partialtrees) == 1
-        return partialtrees[1]
+# A contraction tree for a network of `n` tensors is represented as a list of pairwise
+# contractions `steps`, where `steps[s] = (i, j)` contracts the (partial) results `i` and `j`.
+# Here, results `1:n` are the tensors themselves, and result `n + s` is the result of `steps[s]`.
+# The partial tree containing tensor `k` is `partialtrees[group[k]]`.
+
+# Merge the partial trees containing tensors `k1` and `k2`. The merged tree replaces the one
+# containing the tensor that comes first, and nothing happens if both tensors are already in
+# the same partial tree.
+function _mergetrees!(
+        steps::Vector{NTuple{2, Int}}, partialtrees::Vector{Int}, group::Vector{Int},
+        k1::Int, k2::Int
+    )
+    g1, g2 = minmax(group[k1], group[k2])
+    g1 == g2 && return steps
+    push!(steps, (partialtrees[g1], partialtrees[g2]))
+    partialtrees[g1] = length(group) + length(steps)
+    replace!(group, g2 => g1)
+    return steps
+end
+
+# Combine the remaining partial trees, which are disconnected and thus require outer products,
+# starting from the end.
+function _finalizetree!(
+        steps::Vector{NTuple{2, Int}}, partialtrees::Vector{Int}, group::Vector{Int}
+    )
+    n = length(group)
+    remaining = unique!(group)
+    tree = partialtrees[remaining[end]]
+    for g in Iterators.drop(Iterators.reverse(remaining), 1)
+        push!(steps, (partialtrees[g], tree))
+        tree = n + length(steps)
     end
-    if all(isempty, contractionindices) # disconnected network
-        partialtrees[end - 1] = Any[partialtrees[end - 1], partialtrees[end]]
-        pop!(partialtrees)
-        pop!(contractionindices)
-    else
-        let firstind = minimum(vcat(contractionindices...))
-            i1 = findfirst(x -> in(firstind, x), contractionindices)
-            i2 = findnext(x -> in(firstind, x), contractionindices, i1 + 1)
-            @assert i1 !== nothing && i2 !== nothing
-            newindices = unique2(vcat(contractionindices[i1], contractionindices[i2]))
-            newtree = Any[partialtrees[i1], partialtrees[i2]]
-            partialtrees[i1] = newtree
-            deleteat!(partialtrees, i2)
-            contractionindices[i1] = newindices
-            deleteat!(contractionindices, i2)
-        end
+    return steps
+end
+
+# convert a list of pairwise contractions to a nested contraction tree
+function _nestedtree(steps::Vector{NTuple{2, Int}}, n::Int)
+    trees = collect(Any, 1:n)
+    for (i, j) in steps
+        push!(trees, Any[trees[i], trees[j]])
     end
-    return _ncontree!(partialtrees, contractionindices)
+    return trees[end]
 end
 
 """

From 4044708d297cbf545208326a7ed589253b702fba Mon Sep 17 00:00:00 2001
From: lkdvos <ldevos98@gmail.com>
Date: Mon, 12 Oct 2026 16:06:05 -0400
Subject: [PATCH 4/7] Reduce the overhead of `ncon`

The network and the intermediate results of `ncon` have non-uniform types, which previously
caused type instabilities and runtime type inference throughout the evaluation of the
contraction tree. Now:

- `resolve_traces` canonicalizes the tensors and labels to a `Vector{Any}` and
  `Vector{Vector{Int}}`, and the output and conjugation flags are canonicalized as well,
  such that the remainder of `ncon` is type stable and compiled only once.
- The contraction tree is traversed as a list of pairwise contractions with integer ids.
- The keyword arguments are processed once, instead of for every pairwise contraction.
- All label manipulations happen on `Vector{Int}`, and every pairwise contraction goes
  through a single function barrier, in which the scalar type is inferred statically
  instead of through `promote_op` at runtime.

As a side effect, `ncon` can now be differentiated with Zygote, and networks with traces
can be specified with row matrices as index lists.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 docs/src/changelog.md                    |   4 +
 ext/TensorOperationsChainRulesCoreExt.jl |   7 ++
 src/implementation/ncon.jl               | 109 ++++++++++++++++-------
 src/indexnotation/ncontree.jl            |  23 ++++-
 4 files changed, 108 insertions(+), 35 deletions(-)

diff --git a/docs/src/changelog.md b/docs/src/changelog.md
index 6bd85d79..95f84bc6 100644
--- a/docs/src/changelog.md
+++ b/docs/src/changelog.md
@@ -17,8 +17,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- `ncon` can be differentiated with Zygote, including networks with traces.
+- `ncon` accepts traces in networks whose index lists are given as row matrices, e.g. `[1 1 -1]`.
+
 ### Performance
 
+- `ncon` has significantly less overhead: the network analysis and contraction tree construction are type stable, and every pairwise contraction goes through a single dynamic dispatch.
 - Contractions of `Array`s that require a permuted copy of an input no longer suffer from a type instability.
 - `ncon` partitions its intermediates as `(open indices of A; open indices of B)` instead of placing all indices in the codomain, which avoids repartitioning copies for tensor types such as `TensorMap`s with symmetries ([#311](https://github.com/QuantumKitHub/TensorOperations.jl/issues/311)).
 
diff --git a/ext/TensorOperationsChainRulesCoreExt.jl b/ext/TensorOperationsChainRulesCoreExt.jl
index 1ca4c721..eda32a19 100644
--- a/ext/TensorOperationsChainRulesCoreExt.jl
+++ b/ext/TensorOperationsChainRulesCoreExt.jl
@@ -193,5 +193,12 @@ end
 @non_differentiable TensorOperations.nconoutput(args...)
 @non_differentiable TensorOperations.check_nconstyle(args...)
 @non_differentiable TensorOperations.indexordertree(args...)
+@non_differentiable TensorOperations._ncontree_steps(args...)
+@non_differentiable TensorOperations._indexordertree_steps(args...)
+# label manipulations that mutate arrays, which is not supported by e.g. Zygote
+@non_differentiable TensorOperations._ncon_network(args...)
+@non_differentiable TensorOperations._ncon_labelvector(args...)
+@non_differentiable TensorOperations._ncon_conjlist(args...)
+@non_differentiable TensorOperations._ncon_symdiff(args...)
 
 end
diff --git a/src/implementation/ncon.jl b/src/implementation/ncon.jl
index 79f06df3..24252a38 100644
--- a/src/implementation/ncon.jl
+++ b/src/implementation/ncon.jl
@@ -39,45 +39,90 @@ function ncon(
         end
     end
 
-    tensors, network = resolve_traces(tensors, network)
-    tree = order === nothing ? ncontree(network) : indexordertree(network, order)
+    tensors′, network′ = resolve_traces(tensors, network)
+    # extract the keyword arguments once, rather than for every pairwise contraction
+    backend, allocator = _ncon_kwargs(; kwargs...)
+    conjlist′ = _ncon_conjlist(conjlist)
+    return _ncon(tensors′, network′, conjlist′, output′, order, backend, allocator)
+end
 
-    allocator = haskey(kwargs, :allocator) ? kwargs[:allocator] : DefaultAllocator()
+function _ncon(tensors, network, conjlist, output, order, backend, allocator)
+    steps = order === nothing ? _ncontree_steps(network) :
+        _indexordertree_steps(network, order)
     cp = allocator_checkpoint!(allocator)
-
-    A, IA, conjA = contracttree(tensors, network, conjlist, tree[1]; kwargs...)
-    B, IB, conjB = contracttree(tensors, network, conjlist, tree[2]; kwargs...)
-    IC = tuple(output′...)
-    C = tensorcontract(IC, A, IA, conjA, B, IB, conjB; kwargs...)
-    tree[1] isa Int || tensorfree!(A, allocator)
-    tree[2] isa Int || tensorfree!(B, allocator)
+    n = length(tensors)
+    i, j = steps[end]
+    A, IA, conjA = contracttree(tensors, network, conjlist, steps, i, backend, allocator)
+    B, IB, conjB = contracttree(tensors, network, conjlist, steps, j, backend, allocator)
+    C = _ncon_contract(A, IA, conjA, B, IB, conjB, output, Val(false), backend, allocator)
+    i > n && tensorfree!(A, allocator)
+    j > n && tensorfree!(B, allocator)
     allocator_reset!(allocator, cp)
-    return length(IC) == 0 ? tensorscalar(C) : C
+    return isempty(output) ? tensorscalar(C) : C
+end
+
+# Compute the (partial) result `k` of the contraction tree `steps` (see `_mergetrees!`).
+# The types and number of indices of the intermediate tensors are only known at runtime, but
+# the labels are kept as `Vector{Int}` such that all index computations are type stable, and
+# every pairwise contraction goes through a single dynamic dispatch, in `_ncon_contract`.
+function contracttree(tensors, network, conjlist, steps, k::Int, backend, allocator)
+    n = length(tensors)
+    k <= n && return tensors[k], network[k], conjlist[k]
+    i, j = steps[k - n]
+    A, IA, conjA = contracttree(tensors, network, conjlist, steps, i, backend, allocator)
+    B, IB, conjB = contracttree(tensors, network, conjlist, steps, j, backend, allocator)
+    IC = _ncon_symdiff(IA, IB)
+    C = _ncon_contract(A, IA, conjA, B, IB, conjB, IC, Val(true), backend, allocator)
+    i > n && tensorfree!(A, allocator)
+    j > n && tensorfree!(B, allocator)
+    return C, IC, false
 end
 
-function contracttree(tensors, network, conjlist, tree; kwargs...)
-    @nospecialize
-    if tree isa Int
-        return tensors[tree], tuple(network[tree]...), (conjlist[tree])
+function _ncon_kwargs(; kwargs...)
+    args = _kwargs2args(; kwargs...)
+    return length(args) == 2 ? args :
+        length(args) == 1 ? (args[1], DefaultAllocator()) :
+        (DefaultBackend(), DefaultAllocator())
+end
+
+# labels of the result of contracting tensors with labels `IA` and `IB`, i.e. all labels that
+# appear in only one of them, in order of appearance
+# note: this is much faster than `Base.symdiff`, which constructs a `Set`
+function _ncon_symdiff(IA::Vector{Int}, IB::Vector{Int})
+    IC = Int[]
+    for i in IA
+        i in IB || push!(IC, i)
     end
-    A, IA, conjA = contracttree(tensors, network, conjlist, tree[1]; kwargs...)
-    B, IB, conjB = contracttree(tensors, network, conjlist, tree[2]; kwargs...)
-    IC = tuple(symdiff(IA, IB)...)
-    pA, pB, _ = contract_indices(IA, IB, IC)
-    # `IC` lists the open indices of `A` before those of `B`, as in the product
-    pAB = trivialpermutation(length(pA[1]), length(pB[2]))
-    TC = promote_contract(scalartype(A), scalartype(B))
-    allocator = haskey(kwargs, :allocator) ? kwargs[:allocator] : DefaultAllocator()
-    backend = haskey(kwargs, :backend) ? kwargs[:backend] : DefaultBackend()
-    C = tensoralloc_contract(TC, A, pA, conjA, B, pB, conjB, pAB, Val(true), allocator)
-    C = tensorcontract!(
-        C, A, pA, conjA, B, pB, conjB, pAB, One(), Zero(), backend, allocator
+    for i in IB
+        i in IA || push!(IC, i)
+    end
+    return IC
+end
+
+function _ncon_contract(
+        A, IA::Vector{Int}, conjA::Bool, B, IB::Vector{Int}, conjB::Bool, IC::Vector{Int},
+        istemp::Val, backend, allocator
     )
-    tree[1] isa Int || tensorfree!(A, allocator)
-    tree[2] isa Int || tensorfree!(B, allocator)
-    return C, IC, false
+    # function barrier: the number of indices of `A`, `B` and `C` is only known at runtime
+    IA′, IB′, IC′ = _ncon_tuple(IA), _ncon_tuple(IB), _ncon_tuple(IC)
+    return _ncon_contract!(A, IA′, conjA, B, IB′, conjB, IC′, istemp, backend, allocator)
+end
+function _ncon_contract!(A, IA, conjA, B, IB, conjB, IC, istemp::Val, backend, allocator)
+    pA, pB, pAB = contract_indices(IA, IB, IC)
+    if istemp === Val(true)
+        # intermediates list the open indices of `A` before those of `B`, so partition them as
+        # `(open A; open B)`, which avoids repartitioning copies for e.g. symmetric tensors
+        pAB = trivialpermutation(numout(pA), numin(pB))
+    end
+    TC = promote_contract(scalartype(A), scalartype(B))
+    C = tensoralloc_contract(TC, A, pA, conjA, B, pB, conjB, pAB, istemp, allocator)
+    return tensorcontract!(C, A, pA, conjA, B, pB, conjB, pAB, One(), Zero(), backend, allocator)
 end
 
+# convert a list of labels to a tuple
+# note: `ntuple` is much faster than `Tuple(I)` or `(I...,)` for short lists
+_ncon_tuple(I::Vector{Int}) = ntuple(k -> I[k], length(I))
+
 function nconoutput(network, output)
     outputindices = Vector{Int}()
     for n in network
@@ -91,5 +136,7 @@ function nconoutput(network, output)
 
     issetequal(output, outputindices) ||
         throw(ArgumentError("invalid NCON network: $network -> $output"))
-    return output
+    return collect(Int, output)
 end
+
+_ncon_conjlist(conjlist) = collect(Bool, conjlist)
diff --git a/src/indexnotation/ncontree.jl b/src/indexnotation/ncontree.jl
index b16cee58..be93f042 100644
--- a/src/indexnotation/ncontree.jl
+++ b/src/indexnotation/ncontree.jl
@@ -141,10 +141,25 @@ function nconindexcompletion(ex)
     end
 end
 
+# Resolve the traces in the network, and canonicalize the types of the tensor list and the
+# network, such that the remainder of `ncon` is type stable and only compiled once.
 function resolve_traces(tensors, network)
-    transformed = map(zip(tensors, network)) do (A, IA)
-        IC = unique2(IA)
-        return length(IC) == length(IA) ? (A, IA) : (tensortrace(IC, A, IA), IC)
+    network′ = _ncon_network(network)
+    all(k -> length(network[k]) == length(network′[k]), eachindex(network)) &&
+        return _anyvector(tensors), network′
+    tensors′ = map(tensors, network, network′) do A, IA, IC
+        return length(IA) == length(IC) ? A : tensortrace(IC, A, _ncon_labelvector(IA))
     end
-    return first.(transformed), last.(transformed)
+    return _anyvector(tensors′), network′
 end
+
+# labels after resolving the traces
+_ncon_network(network) =
+    Vector{Int}[_ncon_labelvector(allunique(IA) ? IA : unique2(IA)) for IA in network]
+
+# canonicalize any list of labels, e.g. a `Tuple` or a row matrix like `[1 2 -1]`
+_ncon_labelvector(IA) = vec(collect(Int, IA))
+
+# note: avoiding `collect(Any, xs)` and mutation for compatibility with Zygote
+_anyvector(xs::AbstractVector) = Vector{Any}(xs)
+_anyvector(xs::Tuple) = Any[xs...]

From 82c8936c917fe25599d189139b051fc756904801 Mon Sep 17 00:00:00 2001
From: lkdvos <ldevos98@gmail.com>
Date: Tue, 13 Oct 2026 10:28:02 -0400
Subject: [PATCH 5/7] Add `contract_indices` for labels stored in vectors

The tuple implementation of `contract_indices` manipulates the labels in a type-unstable
manner, and is compiled for every combination of numbers of indices. For labels stored in
vectors, as in `ncon`, the permutations are now computed as vectors in a type-stable manner,
and only converted to tuples at the end.

`ncon` passes these tuples to its function barrier as separate arguments rather than as
`Index2Tuple`s, which avoids boxing them. This reduces the compilation time for networks with
different numbers of indices, at a small cost in runtime for small networks.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 docs/src/changelog.md                    |  1 +
 ext/TensorOperationsChainRulesCoreExt.jl |  1 +
 src/implementation/indices.jl            | 55 +++++++++++++++++++++---
 src/implementation/ncon.jl               | 30 +++++++------
 test/auxiliary.jl                        | 34 +++++++++++++++
 5 files changed, 102 insertions(+), 19 deletions(-)

diff --git a/docs/src/changelog.md b/docs/src/changelog.md
index 95f84bc6..b45884b6 100644
--- a/docs/src/changelog.md
+++ b/docs/src/changelog.md
@@ -24,6 +24,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - `ncon` has significantly less overhead: the network analysis and contraction tree construction are type stable, and every pairwise contraction goes through a single dynamic dispatch.
 - Contractions of `Array`s that require a permuted copy of an input no longer suffer from a type instability.
+- Index labels given as vectors, as in `ncon` or the method syntax `tensorcontract(IC, A, IA, B, IB)`, are converted to index permutations in a type-stable manner, which avoids compiling the conversion for every combination of numbers of indices.
 - `ncon` partitions its intermediates as `(open indices of A; open indices of B)` instead of placing all indices in the codomain, which avoids repartitioning copies for tensor types such as `TensorMap`s with symmetries ([#311](https://github.com/QuantumKitHub/TensorOperations.jl/issues/311)).
 
 ## [5.8.2](https://github.com/QuantumKitHub/TensorOperations.jl/compare/v5.8.1...v5.8.2) - 2026-09-30
diff --git a/ext/TensorOperationsChainRulesCoreExt.jl b/ext/TensorOperationsChainRulesCoreExt.jl
index eda32a19..311f0345 100644
--- a/ext/TensorOperationsChainRulesCoreExt.jl
+++ b/ext/TensorOperationsChainRulesCoreExt.jl
@@ -23,6 +23,7 @@ using LinearAlgebra
 @non_differentiable TensorOperations.tensoralloc_contract(args...)
 @non_differentiable TensorOperations.promote_contract(args...)
 @non_differentiable TensorOperations.promote_add(args...)
+@non_differentiable TensorOperations.contract_indices(args...)
 
 # Cannot free intermediate tensors when using AD
 # Thus we change the forward passes: `istemp=false` and `tensorfree!` is a no-op
diff --git a/src/implementation/indices.jl b/src/implementation/indices.jl
index 06ddb6c0..bafff825 100644
--- a/src/implementation/indices.jl
+++ b/src/implementation/indices.jl
@@ -81,16 +81,13 @@ function trace_indices(IA::NTuple{NA, Any}, IC::NTuple{NC, Any}) where {NA, NC}
     return (indCinA, ()), (cindA1, cindA2)
 end
 
-function contract_indices(IA, IB, IC)
-    return contract_indices(tuple(IA...), tuple(IB...), tuple(IC...))
-end
+contract_indices(IA, IB, IC) = contract_indices(tuple(IA...), tuple(IB...), tuple(IC...))
 function contract_indices(
         IA::NTuple{NA, Any}, IB::NTuple{NB, Any}, IC::NTuple{NC, Any}
     ) where {NA, NB, NC}
     # labels
     IAB = (IA..., IB...)
-    isodd(length(IAB) - length(IC)) &&
-        throw(IndexError("invalid contraction pattern: $IA and $IB to $IC"))
+    isodd(length(IAB) - length(IC)) && _throw_invalid_contraction(IA, IB, IC)
     Icontract = tunique(tsetdiff(IAB, IC))
     IopenA = tsetdiff(IA, Icontract)
     IopenB = tsetdiff(IB, Icontract)
@@ -107,12 +104,58 @@ function contract_indices(
     pAB = (indCinoAB, ())
 
     if !isperm(linearize(pA)) || !isperm(linearize(pB)) || !isperm(indCinoAB)
-        throw(IndexError("invalid contraction pattern: $IA and $IB to $IC"))
+        _throw_invalid_contraction(IA, IB, IC)
     end
 
     return pA, pB, pAB
 end
 
+# Implementation for labels stored in vectors, where the number of labels is only known at
+# runtime. The permutations are computed as vectors in a type stable manner, and only then
+# converted to tuples. This avoids the type instabilities of the tuple implementation, as well
+# as compiling it for every combination of lengths.
+# note: inlined such that callers with a function barrier, such as `ncon`, can pass the
+# resulting tuples on without boxing them
+@inline function contract_indices(IA::AbstractVector, IB::AbstractVector, IC::AbstractVector)
+    # split the indices of `A` into open and contracted ones
+    oindA, cindA, cindB = Int[], Int[], Int[]
+    for (i, l) in enumerate(IA)
+        j = _findfirst(isequal(l), IB)
+        if j == 0
+            push!(oindA, i)
+        else
+            push!(cindA, i)
+            push!(cindB, j)
+        end
+    end
+    oindB = findall(l -> !(l in IA), IB)
+    # position of the indices of `C` in the open indices of `A` and `B`
+    indCinoAB = map(IC) do l
+        i = findfirst(i -> isequal(IA[i], l), oindA)
+        isnothing(i) || return i
+        j = findfirst(j -> isequal(IB[j], l), oindB)
+        return isnothing(j) ? 0 : length(oindA) + j
+    end
+    # every index of `B` and `C` should be accounted for exactly once
+    (
+        length(cindB) + length(oindB) == length(IB) && allunique(cindB) &&
+            length(oindA) + length(oindB) == length(IC) && allunique(indCinoAB) &&
+            !(0 in indCinoAB)
+    ) || throw(IndexError("invalid contraction pattern: $IA and $IB to $IC"))
+
+    pA = (_totuple(oindA), _totuple(cindA))
+    pB = (_totuple(cindB), _totuple(oindB))
+    pAB = (_totuple(indCinoAB), ())
+    return pA, pB, pAB
+end
+
+# note: `ntuple` is much faster than `Tuple` or splatting for short lists, and the separate
+# function avoids inferring the resulting union of tuple types in the caller
+@noinline _totuple(p::Vector{Int}) = ntuple(k -> p[k], length(p))
+
+@noinline _throw_invalid_contraction(IA, IB, IC) =
+    throw(IndexError("invalid contraction pattern: $IA and $IB to $IC"))
+
 #-------------------------------------------------------------------------------------------
 # Convert indices to einsum labels: useful for package extensions / add-ons (e.g. TBLIS)
 #-------------------------------------------------------------------------------------------
diff --git a/src/implementation/ncon.jl b/src/implementation/ncon.jl
index 24252a38..03d688a2 100644
--- a/src/implementation/ncon.jl
+++ b/src/implementation/ncon.jl
@@ -103,26 +103,30 @@ function _ncon_contract(
         A, IA::Vector{Int}, conjA::Bool, B, IB::Vector{Int}, conjB::Bool, IC::Vector{Int},
         istemp::Val, backend, allocator
     )
+    # only `_ncon_contract!` should be specialized on the types of the tensors
+    @nospecialize A B
+    pA, pB, pAB = contract_indices(IA, IB, IC)
     # function barrier: the number of indices of `A`, `B` and `C` is only known at runtime
-    IA′, IB′, IC′ = _ncon_tuple(IA), _ncon_tuple(IB), _ncon_tuple(IC)
-    return _ncon_contract!(A, IA′, conjA, B, IB′, conjB, IC′, istemp, backend, allocator)
+    # note: the tuples are passed separately rather than as `Index2Tuple`s to avoid boxing
+    return _ncon_contract!(
+        A, pA[1], pA[2], conjA, B, pB[1], pB[2], conjB, pAB[1], istemp, backend, allocator
+    )
 end
-function _ncon_contract!(A, IA, conjA, B, IB, conjB, IC, istemp::Val, backend, allocator)
-    pA, pB, pAB = contract_indices(IA, IB, IC)
-    if istemp === Val(true)
-        # intermediates list the open indices of `A` before those of `B`, so partition them as
-        # `(open A; open B)`, which avoids repartitioning copies for e.g. symmetric tensors
-        pAB = trivialpermutation(numout(pA), numin(pB))
-    end
+function _ncon_contract!(
+        A, oindA, cindA, conjA, B, cindB, oindB, conjB, indCinoAB, istemp::Val,
+        backend, allocator
+    )
+    pA = (oindA, cindA)
+    pB = (cindB, oindB)
+    # intermediates list the open indices of `A` before those of `B`, so partition them as
+    # `(open A; open B)`, which avoids repartitioning copies for e.g. symmetric tensors
+    pAB = istemp === Val(true) ? trivialpermutation(length(oindA), length(oindB)) :
+        (indCinoAB, ())
     TC = promote_contract(scalartype(A), scalartype(B))
     C = tensoralloc_contract(TC, A, pA, conjA, B, pB, conjB, pAB, istemp, allocator)
     return tensorcontract!(C, A, pA, conjA, B, pB, conjB, pAB, One(), Zero(), backend, allocator)
 end
 
-# convert a list of labels to a tuple
-# note: `ntuple` is much faster than `Tuple(I)` or `(I...,)` for short lists
-_ncon_tuple(I::Vector{Int}) = ntuple(k -> I[k], length(I))
-
 function nconoutput(network, output)
     outputindices = Vector{Int}()
     for n in network
diff --git a/test/auxiliary.jl b/test/auxiliary.jl
index fcfa95af..e102ca03 100644
--- a/test/auxiliary.jl
+++ b/test/auxiliary.jl
@@ -336,5 +336,39 @@ end
             )
             @test contract_indices(contract_labels(pA, pB, pAB)...) == (pA, pB, pAB)
         end
+
+        # labels stored in vectors
+        @test contract_indices([:a, :b, :c], [:c, :d], [:a, :b, :d]) ==
+            (((1, 2), (3,)), ((1,), (2,)), ((1, 2, 3), ()))
+        @test contract_indices([:c, :a, :b], [:c, :d], [:a, :b, :d]) ==
+            (((2, 3), (1,)), ((1,), (2,)), ((1, 2, 3), ()))
+        @test contract_indices([1, 2], [3, 4], [3, 1, 4, 2]) ==
+            (((1, 2), ()), ((), (1, 2)), ((3, 1, 4, 2), ()))
+        @test contract_indices([1, 2], [2, 1], Int[]) ==
+            (((), (1, 2)), ((2, 1), ()), ((), ()))
+        @test_throws IndexError contract_indices([:a, :b, :c], [:c, :d], [:a, :b, :d, :e])
+        @test_throws IndexError contract_indices([:a, :b, :c], [:c, :d], [:a, :b, :e])
+        @test_throws IndexError contract_indices([:a, :b, :c], [:c, :d, :e], [:a, :b, :d])
+        @test_throws IndexError contract_indices([:a, :b, :c], [:c, :c], [:a, :b])
+
+        # agrees with the tuple implementation, including for invalid patterns
+        result(f) = try
+            f()
+        catch e
+            e isa IndexError || rethrow()
+            IndexError
+        end
+        for _ in 1:1000
+            c = collect(1:rand(0:3))
+            IA = shuffle([c; 100 .+ (1:rand(0:3))])
+            IB = shuffle([c; 200 .+ (1:rand(0:3))])
+            IC = shuffle(filter(>(99), [IA; IB]))
+            r = rand(1:4)
+            r == 1 && !isempty(IC) && (IC[1] = 999)
+            r == 2 && !isempty(IA) && (IA[1] = 777)
+            r == 3 && length(IB) > 1 && (IB[2] = IB[1])
+            @test result(() -> contract_indices(IA, IB, IC)) ==
+                result(() -> contract_indices(Tuple(IA), Tuple(IB), Tuple(IC)))
+        end
     end
 end

From 600aa4345ca135415dfeaaa6483815ad84d472d9 Mon Sep 17 00:00:00 2001
From: lkdvos <ldevos98@gmail.com>
Date: Tue, 13 Oct 2026 13:00:28 -0400
Subject: [PATCH 6/7] Avoid boxing the BLAS-contractable inputs in
 `_blas_contract!`

`makeblascontractable` returned either its input or a permuted copy, whose types differ when
the copy is allocated by an allocator such as Bumper's. The resulting union is not an
`isbits` type and was therefore heap allocated on every contraction, even when no copy was
required. Its specialization on the scalar type also made Enzyme fail ("could not determine
how to keep a pointer phi rooted").

`_blas_contract!` now makes the inputs BLAS-contractable one at a time, passing them on to
the next step in each branch rather than returning them, such that the union is never
materialized. Each input is reshaped into a matrix before it is passed on, so that the
subsequent steps no longer depend on the number of indices of the inputs, which keeps the
number of compiled methods (and the compilation time) at the level of before this change.
`makeblascontractable` and `_unsafe_blas_contract!` are no longer used and are removed,
together with the tests of the former.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 docs/src/changelog.md              |  1 +
 src/implementation/blascontract.jl | 95 ++++++++++++------------------
 test/tensor.jl                     | 18 +-----
 3 files changed, 41 insertions(+), 73 deletions(-)

diff --git a/docs/src/changelog.md b/docs/src/changelog.md
index b45884b6..399232cc 100644
--- a/docs/src/changelog.md
+++ b/docs/src/changelog.md
@@ -19,6 +19,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - `ncon` can be differentiated with Zygote, including networks with traces.
 - `ncon` accepts traces in networks whose index lists are given as row matrices, e.g. `[1 1 -1]`.
+- Contractions using the `StridedBLAS` backend no longer heap allocate a view of an input tensor when an allocator such as Bumper's is used ([#273](https://github.com/QuantumKitHub/TensorOperations.jl/issues/273)).
 
 ### Performance
 
diff --git a/src/implementation/blascontract.jl b/src/implementation/blascontract.jl
index 289ae0c2..c0e1325e 100644
--- a/src/implementation/blascontract.jl
+++ b/src/implementation/blascontract.jl
@@ -41,74 +41,57 @@ function blas_contract!(
     return C
 end
 
-# implement necessary permutations
+# implement necessary permutations - each input is made BLAS-contractable and reshaped into a
+# matrix in a separate step within each branch, to avoid boxing/type instabilities from
+# returning a union of the input and its permuted copy; as the subsequent steps only receive
+# matrices, they no longer depend on the number of indices of the inputs
 function _blas_contract!(C, A, pA, B, pB, pAB, α, β, backend, allocator)
-    TC = eltype(C)
-
-    A_, pA, flagA = makeblascontractable(A, pA, TC, backend, allocator)
-    B_, pB, flagB = makeblascontractable(B, pB, TC, backend, allocator)
-
-    # the partition of `ipAB` is required by `isblasdestination` and `tensoralloc_add`, but
-    # the actual contraction only needs the linearized permutation
     ipAB = oindABinC(pAB, pA, pB)
-    ipAB′ = linearize(ipAB)
-    flagC = isblasdestination(C, ipAB)
-    if flagC
-        C_ = C
-        _unsafe_blas_contract!(C_, A_, pA, B_, pB, ipAB′, α, β)
+    if isblascontractable(A, pA) && eltype(A) == eltype(C)
+        return _blas_contract_B!(C, _blas_matrix(A, pA), B, pB, pAB, ipAB, α, β, backend, allocator)
+    else
+        A_ = tensoralloc_add(eltype(C), A, pA, false, Val(true), allocator)
+        Anew = tensoradd!(SV(A_), A, pA, false, One(), Zero(), backend, allocator)
+        Amat = _blas_matrix(Anew, trivialpermutation(pA))
+        _blas_contract_B!(C, Amat, B, pB, pAB, ipAB, α, β, backend, allocator)
+        tensorfree!(A_, allocator)
+        return C
+    end
+end
+function _blas_contract_B!(C, A, B, pB, pAB, ipAB, α, β, backend, allocator)
+    if isblascontractable(B, pB) && eltype(B) == eltype(C)
+        return _blas_contract_C!(C, A, _blas_matrix(B, pB), pAB, ipAB, α, β, backend, allocator)
+    else
+        B_ = tensoralloc_add(eltype(C), B, pB, false, Val(true), allocator)
+        Bnew = tensoradd!(SV(B_), B, pB, false, One(), Zero(), backend, allocator)
+        Bmat = _blas_matrix(Bnew, trivialpermutation(pB))
+        _blas_contract_C!(C, A, Bmat, pAB, ipAB, α, β, backend, allocator)
+        tensorfree!(B_, allocator)
+        return C
+    end
+end
+function _blas_contract_C!(C, A, B, pAB, ipAB, α, β, backend, allocator)
+    if isblasdestination(C, ipAB)
+        mul!(_blas_matrix(C, ipAB), A, B, α, β)
     else
+        TC = eltype(C)
         C_ = SV(tensoralloc_add(TC, C, ipAB, false, Val(true), allocator))
-        _unsafe_blas_contract!(
-            C_, A_, pA, B_, pB, trivialpermutation(ipAB′),
-            one(TC), zero(TC)
-        )
+        mul!(_blas_matrix(C_, trivialpermutation(ipAB)), A, B, one(TC), zero(TC))
         # `C` already exists, so only the permutation of `pAB` matters here and it can be
         # handed over in the canonical `{N,0}` partition
         tensoradd!(C, C_, (pAB, ()), false, α, β, backend, allocator)
         tensorfree!(C_.parent, allocator)
     end
-    flagA || tensorfree!(A_.parent, allocator)
-    flagB || tensorfree!(B_.parent, allocator)
     return C
 end
 
-# perform the actual contraction, assuming it can be done as matrix multiplication by simply
-# reshaping without any further allocations
-function _unsafe_blas_contract!(
-        C::StridedView{T},
-        A::StridedView{T}, pA,
-        B::StridedView{T}, pB,
-        pAB::IndexTuple, α, β
-    ) where {T <: BlasFloat}
-    sizeA = size(A)
-    sizeB = size(B)
-    csizeA = TupleTools.getindices(sizeA, pA[2])
-    csizeB = TupleTools.getindices(sizeB, pB[1])
-    osizeA = TupleTools.getindices(sizeA, pA[1])
-    osizeB = TupleTools.getindices(sizeB, pB[2])
-
-    mul!(
-        sreshape(permutedims(C, pAB), (prod(osizeA), prod(osizeB))),
-        sreshape(permutedims(A, linearize(pA)), (prod(osizeA), prod(csizeA))),
-        sreshape(permutedims(B, linearize(pB)), (prod(csizeB), prod(osizeB))),
-        α, β
-    )
-
-    return C
-end
-
-function makeblascontractable(A, pA, ::Type{TC}, backend, allocator) where {TC}
-    flagA = isblascontractable(A, pA) && eltype(A) == TC
-    if !flagA
-        A_ = tensoralloc_add(TC, A, pA, false, Val(true), allocator)
-        Anew = SV(A_)
-        Anew = tensoradd!(Anew, A, pA, false, One(), Zero(), backend, allocator)
-        pAnew = trivialpermutation(pA)
-    else
-        Anew = A
-        pAnew = pA
-    end
-    return Anew, pAnew, flagA
+# reshape `A` into the matrix with the indices `p[1]` as rows and `p[2]` as columns, assuming
+# this is possible without any further allocations
+function _blas_matrix(A::StridedView, p::Index2Tuple)
+    sz = size(A)
+    m = prod(TupleTools.getindices(sz, p[1]))
+    n = prod(TupleTools.getindices(sz, p[2]))
+    return sreshape(permutedims(A, linearize(p)), (m, n))
 end
 
 function isblascontractable(A::StridedView, p::Index2Tuple)
diff --git a/test/tensor.jl b/test/tensor.jl
index 2ac3f71e..11ed2473 100644
--- a/test/tensor.jl
+++ b/test/tensor.jl
@@ -543,40 +543,24 @@ end
     end
 
     @testset "methods for StridedBLAS" begin
-        using TensorOperations: isblascontractable, isblasdestination, makeblascontractable,
-            StridedBLAS, DefaultAllocator
+        using TensorOperations: isblascontractable
         using Strided
-        backend = StridedBLAS()
-        allocator = DefaultAllocator()
 
         A = StridedView(rand(ComplexF64, 4, 4, 4, 4))
 
         p = ((1, 2), (3, 4))
         @test isblascontractable(A, p)
-        Anew, pnew, flag = makeblascontractable(A, p, ComplexF64, backend, allocator)
-        @test isblascontractable(Anew, pnew)
-        @test Anew === A
-
         @test !isblascontractable(conj(A), p)
-        Anew, pnew, flag = makeblascontractable(A, p, ComplexF64, backend, allocator)
-        @test isblascontractable(Anew, pnew)
 
         for p in (((2, 1), (3, 4)), ((1,), (3, 2, 4)), ((2, 1, 4), (3,))),
                 op in (identity, conj)
 
             @test !isblascontractable(op(A), p)
-            Anew, pnew, flag = makeblascontractable(
-                op(A), p, ComplexF64, backend, allocator
-            )
-            @test isblascontractable(Anew, pnew)
         end
 
         vA = view(A, 1:2, 1:2, 1:2, 1:2)
         p = ((1, 2), (3, 4))
         @test !isblascontractable(vA, p)
-        Anew, pnew, flag = makeblascontractable(vA, p, ComplexF64, backend, allocator)
-        @test !flag
-        @test isblascontractable(Anew, pnew)
 
         pA = permutedims(A, (3, 4, 1, 2))
         p = ((1, 2), (3, 4))

From 5816040bdc67d2a01bf5e137c0185eecc8efa029 Mon Sep 17 00:00:00 2001
From: lkdvos <ldevos98@gmail.com>
Date: Fri, 23 Oct 2026 09:52:25 -0400
Subject: [PATCH 7/7] Simplify the `ncon` network analysis after review
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- `_ncontree_steps` merges the two occurrences of every contraction index directly from the
  sorted list, rather than searching for the second occurrence.
- `_check_nconstyle` compares neighbours in the sorted list of labels instead of counting.
- `_indexordertree_steps` asserts that the second occurrence of a contraction index exists,
  as before, and uses `∋(c)` rather than anonymous functions.
- Disconnected parts of the network are combined from left to right.
- `resolve_traces` uses comprehensions, without the helpers `_ncon_network` and
  `_anyvector`. As a consequence, `ncon` can no longer be differentiated with Zygote.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 ext/TensorOperationsChainRulesCoreExt.jl |  2 +-
 src/indexnotation/indexordertree.jl      | 10 ++--
 src/indexnotation/ncontree.jl            | 64 +++++++++---------------
 3 files changed, 31 insertions(+), 45 deletions(-)

diff --git a/ext/TensorOperationsChainRulesCoreExt.jl b/ext/TensorOperationsChainRulesCoreExt.jl
index 311f0345..eeee38a7 100644
--- a/ext/TensorOperationsChainRulesCoreExt.jl
+++ b/ext/TensorOperationsChainRulesCoreExt.jl
@@ -197,7 +197,7 @@ end
 @non_differentiable TensorOperations._ncontree_steps(args...)
 @non_differentiable TensorOperations._indexordertree_steps(args...)
 # label manipulations that mutate arrays, which is not supported by e.g. Zygote
-@non_differentiable TensorOperations._ncon_network(args...)
+@non_differentiable TensorOperations.unique2(args...)
 @non_differentiable TensorOperations._ncon_labelvector(args...)
 @non_differentiable TensorOperations._ncon_conjlist(args...)
 @non_differentiable TensorOperations._ncon_symdiff(args...)
diff --git a/src/indexnotation/indexordertree.jl b/src/indexnotation/indexordertree.jl
index 2a87a6af..f8fde981 100644
--- a/src/indexnotation/indexordertree.jl
+++ b/src/indexnotation/indexordertree.jl
@@ -1,4 +1,4 @@
-# verify that every conraction index appears exactly twice, and any other index exactly once
+# verify that every contraction index appears exactly twice, and any other index exactly once
 function _verify_indexorder(network, indexorder)
     for indices in network
         for i in indices
@@ -22,11 +22,11 @@ function _indexordertree_steps(network, indexorder)
     partialtrees = collect(1:length(network))
     group = collect(1:length(network))
     for c in indexorder
-        k1 = findfirst(indices -> c in indices, network)
+        # trace indices have already been removed from the network
+        k1 = findfirst(∋(c), network)
         isnothing(k1) && continue
-        # trace indices have already been removed, so `c` appears in another tensor
-        k2 = findnext(indices -> c in indices, network, k1 + 1)
-        isnothing(k2) && continue
+        k2 = findnext(∋(c), network, k1 + 1)
+        @assert !isnothing(k2) # guaranteed by `_verify_indexorder`
         _mergetrees!(steps, partialtrees, group, k1, k2)
     end
     return _finalizetree!(steps, partialtrees, group)
diff --git a/src/indexnotation/ncontree.jl b/src/indexnotation/ncontree.jl
index be93f042..0b613222 100644
--- a/src/indexnotation/ncontree.jl
+++ b/src/indexnotation/ncontree.jl
@@ -5,33 +5,29 @@ check_nconstyle(network) = _check_nconstyle(network, Val(false))
 function _check_nconstyle(network, ::Val{check}) where {check}
     allindices = Vector{Int}()
     for ind in network
-        for i in ind
-            isa(i, Integer) || return check ? false :
-                throw(IndexError("All indices must be integers"))
-            push!(allindices, i)
-        end
+        all(i -> isa(i, Integer), ind) || return check ? false :
+            throw(IndexError("All indices must be integers"))
+        append!(allindices, ind)
     end
+    # after sorting, equal labels are adjacent
     sort!(allindices)
+    n = length(allindices)
     j = 1
-    while j <= length(allindices)
+    while j <= n
         i = allindices[j]
-        # number of times that `i` appears
-        n = 1
-        while j + n <= length(allindices) && allindices[j + n] == i
-            n += 1
-        end
         if i > 0 # positive labels represent contractions or traces and should appear twice
-            n == 1 && return check ? false :
+            (j == n || allindices[j + 1] != i) && return check ? false :
                 throw(IndexError(lazy"Index $i appears only once in the network"))
-            n > 2 && return check ? false :
+            (j + 2 <= n && allindices[j + 2] == i) && return check ? false :
                 throw(IndexError(lazy"Index $i appears more than twice in the network"))
+            j += 2
         elseif i < 0 # negative labels represent open indices and should appear once
-            n > 1 && return check ? false :
+            (j < n && allindices[j + 1] == i) && return check ? false :
                 throw(IndexError(lazy"Index $i appears more than once in the network"))
+            j += 1
         else # i == 0
             return check ? false : throw(IndexError("Index 0 is not allowed in the network"))
         end
-        j += n
     end
     return check ? true : nothing
 end
@@ -56,10 +52,10 @@ function _ncontree_steps(network)
     steps = NTuple{2, Int}[]
     partialtrees = collect(1:length(network))
     group = collect(1:length(network))
-    for j in sortperm(contractionindices)
-        k = findnext(==(contractionindices[j]), contractionindices, j + 1)
-        # each index appears twice, only consider the first occurrence
-        isnothing(k) || _mergetrees!(steps, partialtrees, group, owners[j], owners[k])
+    # every contraction index appears twice, so after sorting its occurrences are adjacent
+    perm = sortperm(contractionindices)
+    for p in 1:2:length(perm)
+        _mergetrees!(steps, partialtrees, group, owners[perm[p]], owners[perm[p + 1]])
     end
     return _finalizetree!(steps, partialtrees, group)
 end
@@ -84,16 +80,15 @@ function _mergetrees!(
     return steps
 end
 
-# Combine the remaining partial trees, which are disconnected and thus require outer products,
-# starting from the end.
+# Combine the remaining partial trees, which are disconnected and thus require outer products.
 function _finalizetree!(
         steps::Vector{NTuple{2, Int}}, partialtrees::Vector{Int}, group::Vector{Int}
     )
     n = length(group)
     remaining = unique!(group)
-    tree = partialtrees[remaining[end]]
-    for g in Iterators.drop(Iterators.reverse(remaining), 1)
-        push!(steps, (partialtrees[g], tree))
+    tree = partialtrees[remaining[1]]
+    for k in 2:length(remaining)
+        push!(steps, (tree, partialtrees[remaining[k]]))
         tree = n + length(steps)
     end
     return steps
@@ -144,22 +139,13 @@ end
 # Resolve the traces in the network, and canonicalize the types of the tensor list and the
 # network, such that the remainder of `ncon` is type stable and only compiled once.
 function resolve_traces(tensors, network)
-    network′ = _ncon_network(network)
-    all(k -> length(network[k]) == length(network′[k]), eachindex(network)) &&
-        return _anyvector(tensors), network′
-    tensors′ = map(tensors, network, network′) do A, IA, IC
-        return length(IA) == length(IC) ? A : tensortrace(IC, A, _ncon_labelvector(IA))
-    end
-    return _anyvector(tensors′), network′
+    network′ = Vector{Int}[unique2(IA) for IA in network]
+    tensors′ = Any[
+        length(IA) == length(IC) ? A : tensortrace(IC, A, _ncon_labelvector(IA))
+            for (A, IA, IC) in zip(tensors, network, network′)
+    ]
+    return tensors′, network′
 end
 
-# labels after resolving the traces
-_ncon_network(network) =
-    Vector{Int}[_ncon_labelvector(allunique(IA) ? IA : unique2(IA)) for IA in network]
-
 # canonicalize any list of labels, e.g. a `Tuple` or a row matrix like `[1 2 -1]`
 _ncon_labelvector(IA) = vec(collect(Int, IA))
-
-# note: avoiding `collect(Any, xs)` and mutation for compatibility with Zygote
-_anyvector(xs::AbstractVector) = Vector{Any}(xs)
-_anyvector(xs::Tuple) = Any[xs...]