From b93351d710fb33bdcaac891b0ff680db6c0235ea Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Tue, 27 Oct 2026 10:55:56 -0400 Subject: [PATCH 1/2] Avoid hashing in the greedy contraction order Match index names by linear scans instead of `Set`-based `union`, `symdiff` and `isdisjoint`, and update the pending operands in place. The returned tree is unchanged. Co-Authored-By: Claude Opus 5.5 (1M context) --- Project.toml | 2 +- src/prod_tensors.jl | 25 ++++++++++++++++++------- 2 files changed, 19 insertions(+), 8 deletions(-) diff --git a/Project.toml b/Project.toml index c614f208..94e3b7e2 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "ITensorNetworksNext" uuid = "302f2e75-49f0-4526-aef7-d8ba550cb06c" -version = "0.11.1" +version = "0.11.2" authors = ["ITensor developers and contributors"] [workspace] diff --git a/src/prod_tensors.jl b/src/prod_tensors.jl index a9d8a879..acc0dacc 100644 --- a/src/prod_tensors.jl +++ b/src/prod_tensors.jl @@ -1,6 +1,6 @@ using AbstractTrees: nodevalue using Combinatorics: combinations -using ITensorBase: inds +using ITensorBase: inds, name """ ContractionTreeAlgorithm @@ -68,22 +68,33 @@ function contraction_tree(::Greedy, tensors) # The inds of each pending operand. An intermediate never materializes, so its inds are # tracked here rather than read off a tensor. is = [inds(tensors[k]) for k in ks] + # Indices are matched by name with linear scans rather than `Set`s: operands carry only a + # few indices, and hashing an index can be expensive (it may hash a whole graded axis). + # Index equality is also by name (plus the ungraded length), so a bond matches its dual. + hasname(js, i) = any(j -> name(j) == name(i), js) while length(trees) > 1 i1, i2 = argmin(combinations(eachindex(trees), 2)) do (i, j) # Defer outer products: with nothing contracted they cost the full product of both # operands, and taking one early inflates every contraction that follows. - isdisjoint(is[i], is[j]) && return typemax(Int) + any(k -> hasname(is[j], k), is[i]) || return typemax(Int) # Every index on either operand is iterated once, shared ones counted once. - return prod(length, union(is[i], is[j]); init = 1) + unique_j = (k for k in is[j] if !hasname(is[i], k)) + return prod(length, is[i]; init = 1) * prod(length, unique_j; init = 1) end # Shared indices are summed over, so the result carries the symmetric difference. - contracted = symdiff(is[i1], is[i2]) + contracted = [ + filter(k -> !hasname(is[i2], k), is[i1]); + filter(k -> !hasname(is[i1], k), is[i2]) + ] tree = ContractionTree(trees[i1], trees[i2]) # Remove the pair by position. Removing it by value would also drop any other operand # equal to it, silently losing a tensor from a network with repeated tensors. - keep = [i for i in eachindex(trees) if i ∉ (i1, i2)] - trees = [trees[keep]; [tree]] - is = [is[keep]; [contracted]] + # Updating in place rather than rebinding also keeps `is`, captured by the closures above, + # from being boxed. + deleteat!(trees, (i1, i2)) + push!(trees, tree) + deleteat!(is, (i1, i2)) + push!(is, contracted) end return only(trees) end From 04518a883db7535ca7b4ab23e13bf70c5e24eea7 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Tue, 27 Oct 2026 11:22:27 -0400 Subject: [PATCH 2/2] Put the greedy contraction order behind a function barrier Co-Authored-By: Claude Opus 5.5 (1M context) --- src/prod_tensors.jl | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/prod_tensors.jl b/src/prod_tensors.jl index acc0dacc..37cb5ac2 100644 --- a/src/prod_tensors.jl +++ b/src/prod_tensors.jl @@ -64,10 +64,11 @@ struct Greedy <: ContractionTreeAlgorithm end function contraction_tree(::Greedy, tensors) ks = collect(keys(tensors)) isempty(ks) && throw(ArgumentError("No tensors to contract.")) + # Function barrier, since `inds` is not type stable. + return _contraction_tree_greedy(ks, [inds(tensors[k]) for k in ks]) +end +function _contraction_tree_greedy(ks, is) trees = map(ContractionTree, ks) - # The inds of each pending operand. An intermediate never materializes, so its inds are - # tracked here rather than read off a tensor. - is = [inds(tensors[k]) for k in ks] # Indices are matched by name with linear scans rather than `Set`s: operands carry only a # few indices, and hashing an index can be expensive (it may hash a whole graded axis). # Index equality is also by name (plus the ungraded length), so a bond matches its dual.