From ebb00f6a3a9253230d1097e84a3ba98a3dc39303 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Mon, 5 Oct 2026 13:38:52 -0400 Subject: [PATCH 1/5] Represent contraction orders as a ContractionTree Contraction orders become an explicit binary tree over the network's keys, contracted by `prod_tensors(tensors, tree)`, in place of the symbolic expressions built out of ITensorBase's lazy layer. Order finding moves here, its only consumer. A `NormNetwork` is contracted through `flatten_network`, which gives each layer its own vertex so the site index they share is visible to the optimizer. Co-Authored-By: Claude Opus 5 (1M context) --- Project.toml | 12 +- docs/Project.toml | 2 +- examples/Project.toml | 2 +- ...etworksNextOMEinsumContractionOrdersExt.jl | 28 +++- src/ITensorNetworksNext.jl | 2 + src/beliefpropagation/beliefpropagation.jl | 8 +- src/contract_network.jl | 82 ++++-------- src/contraction_tree.jl | 74 +++++++++++ src/normnetwork.jl | 40 +++++- src/prod_tensors.jl | 81 ++++++++++++ test/Project.toml | 6 +- test/test_beliefpropagation.jl | 22 ++- test/test_contract_network.jl | 27 ++-- test/test_contraction_tree.jl | 125 ++++++++++++++++++ test/test_normnetwork.jl | 42 +++++- test/test_tensornetwork.jl | 2 +- 16 files changed, 441 insertions(+), 114 deletions(-) create mode 100644 src/contraction_tree.jl create mode 100644 src/prod_tensors.jl create mode 100644 test/test_contraction_tree.jl diff --git a/Project.toml b/Project.toml index 36c86768..500e476c 100644 --- a/Project.toml +++ b/Project.toml @@ -1,12 +1,13 @@ name = "ITensorNetworksNext" uuid = "302f2e75-49f0-4526-aef7-d8ba550cb06c" -version = "0.10.10" +version = "0.11.0" authors = ["ITensor developers and contributors"] [workspace] projects = ["benchmark", "dev", "docs", "examples", "test"] [deps] +AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c" Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" AlgorithmsInterface = "d1e3940c-cd12-4505-8585-b0a4b322527d" Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa" @@ -22,21 +23,27 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" SimpleTraits = "699a6c99-e7fa-54fc-8d76-47d257e15c1d" SplitApplyCombine = "03a91e81-4c3e-53e1-a0a4-9c0c8f19dd66" TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a" +WrappedUnions = "325db55a-9c6c-5b90-b1a2-ec87e7a38c44" [weakdeps] OMEinsumContractionOrders = "6f22d1fd-8eed-4bb7-9776-e7d684900715" +[sources.ITensorBase] +rev = "mf/remove-lazy" +url = "https://github.com/ITensor/ITensorBase.jl" + [extensions] ITensorNetworksNextOMEinsumContractionOrdersExt = "OMEinsumContractionOrders" [compat] +AbstractTrees = "0.4.5" Adapt = "4.3" AlgorithmsInterface = "0.1" Combinatorics = "1" DataGraphs = "0.6" Dictionaries = "0.4.5" Graphs = "1.13.1" -ITensorBase = "0.15.5" +ITensorBase = "0.16" LinearAlgebra = "1.10" MacroTools = "0.5.16" MatrixAlgebraKit = "0.6" @@ -46,4 +53,5 @@ Random = "1.10" SimpleTraits = "0.9.5" SplitApplyCombine = "1.2.3" TensorAlgebra = "0.23" +WrappedUnions = "0.3" julia = "1.10" diff --git a/docs/Project.toml b/docs/Project.toml index 6adb17f4..eee8bc4e 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -10,5 +10,5 @@ path = ".." [compat] Documenter = "1" ITensorFormatter = "0.2.27" -ITensorNetworksNext = "0.10" +ITensorNetworksNext = "0.11" Literate = "2" diff --git a/examples/Project.toml b/examples/Project.toml index 406118b8..0fa22c56 100644 --- a/examples/Project.toml +++ b/examples/Project.toml @@ -5,4 +5,4 @@ ITensorNetworksNext = "302f2e75-49f0-4526-aef7-d8ba550cb06c" path = ".." [compat] -ITensorNetworksNext = "0.10" +ITensorNetworksNext = "0.11" diff --git a/ext/ITensorNetworksNextOMEinsumContractionOrdersExt/ITensorNetworksNextOMEinsumContractionOrdersExt.jl b/ext/ITensorNetworksNextOMEinsumContractionOrdersExt/ITensorNetworksNextOMEinsumContractionOrdersExt.jl index abede5e5..b968975c 100644 --- a/ext/ITensorNetworksNextOMEinsumContractionOrdersExt/ITensorNetworksNextOMEinsumContractionOrdersExt.jl +++ b/ext/ITensorNetworksNextOMEinsumContractionOrdersExt/ITensorNetworksNextOMEinsumContractionOrdersExt.jl @@ -1,10 +1,30 @@ module ITensorNetworksNextOMEinsumContractionOrdersExt -using ITensorNetworksNext: ITensorNetworksNext -using OMEinsumContractionOrders: CodeOptimizer +using ITensorBase: name +using ITensorNetworksNext: ITensorNetworksNext, ContractionTree +using OMEinsumContractionOrders: + OMEinsumContractionOrders, CodeOptimizer, EinCode, NestedEinsum, optimize_code -function ITensorNetworksNext.contraction_order(alg::CodeOptimizer, tn) - return ITensorNetworksNext._contraction_order(alg, tn) +# Rebuild a `ContractionTree` from an optimized `NestedEinsum`, mapping each leaf's +# `tensorindex` back to the key it came from. +function nested_einsum_to_tree(ks, code::NestedEinsum) + # A leaf holds the 1-based index of its input tensor; internal nodes hold `-1`. + code.tensorindex != -1 && return ContractionTree(ks[code.tensorindex]) + return reduce(ContractionTree, map(Base.Fix1(nested_einsum_to_tree, ks), code.args)) +end + +# Find a contraction order with any OMEinsumContractionOrders optimizer (`GreedyMethod`, +# `TreeSA`, `KaHyParBipartite`, ...) by forwarding to `optimize_code`. +function ITensorNetworksNext.contraction_tree(alg::CodeOptimizer, tensors) + ks = collect(keys(tensors)) + ixs = [map(name, collect(axes(tensors[k]))) for k in ks] + all_inds = reduce(vcat, ixs) + labels = unique(all_inds) + size_dict = Dict(name(ax) => length(ax) for k in ks for ax in axes(tensors[k])) + # Open indices (appearing on a single tensor) are the output of the network. + iy = filter(i -> count(==(i), all_inds) == 1, labels) + code = optimize_code(EinCode(ixs, iy), size_dict, alg) + return nested_einsum_to_tree(ks, code) end end diff --git a/src/ITensorNetworksNext.jl b/src/ITensorNetworksNext.jl index fe1b2540..7025975d 100644 --- a/src/ITensorNetworksNext.jl +++ b/src/ITensorNetworksNext.jl @@ -16,6 +16,8 @@ include("tensornetwork.jl") include("normnetwork.jl") include("normnetworkview.jl") include("ITensorNetworkGenerators/ITensorNetworkGenerators.jl") +include("contraction_tree.jl") +include("prod_tensors.jl") include("contract_network.jl") include("beliefpropagation/messagecache.jl") diff --git a/src/beliefpropagation/beliefpropagation.jl b/src/beliefpropagation/beliefpropagation.jl index 3741fb3a..5a68def3 100644 --- a/src/beliefpropagation/beliefpropagation.jl +++ b/src/beliefpropagation/beliefpropagation.jl @@ -238,10 +238,10 @@ end end # The tensors making up the factor at `vertex`, as separate operands for `contract_network`. A -# `NormNetwork`'s factor is a lazy `ket * conj(bra)` product, and the contraction order sees each -# operand as one node carrying only its outer axes — which hides the physical index the two layers -# share, forcing the doubled vertex to be formed before any message is absorbed (χ^(2 * degree) -# rather than the χ^(degree + 1) an interleaved order reaches). +# `NormNetwork`'s factor is the product of its ket and bra layers. Handing that product over as one +# operand would hide the physical index the two layers share, forcing the doubled vertex to be +# formed before any message is absorbed (χ^(2 * degree) rather than the χ^(degree + 1) an +# interleaved order reaches). factor_tensors(factors, vertex) = [factors[vertex]] function factor_tensors(factors::NormNetwork, vertex) return [kettensor(factors, vertex), bratensor(factors, vertex)] diff --git a/src/contract_network.jl b/src/contract_network.jl index f13b36c8..981338ba 100644 --- a/src/contract_network.jl +++ b/src/contract_network.jl @@ -1,71 +1,35 @@ -using Base.Broadcast: materialize using Base: @kwdef -using ITensorBase: EvaluationOrderAlgorithm, Greedy, Mul, lazy, optimize_evaluation_order, - substitute, symnamedtensor -# `contract_network` -@kwdef struct Exact{Order, OrderAlg} - order::Order = nothing - order_alg::OrderAlg = Greedy() +""" + Exact(; contraction_tree = nothing, tree_alg = Greedy()) + +Contract a network exactly, in the pairwise order given by `contraction_tree`. When no tree is +supplied one is found with `tree_alg`. +""" +@kwdef struct Exact{Tree, TreeAlg} + contraction_tree::Tree = nothing + tree_alg::TreeAlg = Greedy() end function contract_network(alg, tn) return throw(ArgumentError("`contract_network` algorithm `$(alg)` not implemented.")) end -function contract_network(tn; alg = Exact()) - return contract_network(alg, tn) -end +contract_network(tn; alg = Exact()) = contract_network(alg, tn) -# `contract_network(::Exact, ...)` -function get_order(alg::Exact, tn) - # Allow specifying either an explicit `order` or an `order_alg` to compute one. - order = if !isnothing(alg.order) - alg.order - else - contraction_order(tn; alg = alg.order_alg) - end - # Contraction order may or may not have indices attached, canonicalize the format - # by attaching indices. - subs = - Dict(symnamedtensor(i) => symnamedtensor(i, Tuple(axes(t))) for (i, t) in pairs(tn)) - return substitute(order, subs) -end -# Promote the operands to their common type before lowering to the lazy expression, so every lazy -# operand shares one concrete type. Otherwise a network of mixed types (a plain tensor is a trivial -# operator, so mixing operators and plain tensors is the common case) widens the symbolic `Mul` -# container to a `UnionAll` it cannot construct. `promote_type`/`convert` keep an all-plain network -# at the plain type (the promotion is a no-op), so its fast path is unchanged. function contract_network(alg::Exact, tn) - order = get_order(alg, tn) - T = mapreduce(typeof, promote_type, tn) - syms_to_ts = Dict( - symnamedtensor(i, Tuple(axes(t))) => lazy(convert(T, t)) for (i, t) in pairs(tn) - ) - tn_expression = substitute(order, syms_to_ts) - return materialize(tn_expression) + tree = @something alg.contraction_tree contraction_tree(tn; alg = alg.tree_alg) + return prod_tensors(tn, tree) end -# `contraction_order` -function contraction_order end -function contraction_order(tn; alg = Greedy()) - return contraction_order(alg, tn) -end -# Convert the tensor network to a flat symbolic multiplication expression. -struct Flat end -function contraction_order(alg::Flat, tn) - # Same as: `reduce((a, b) -> *(a, b; flatten = true), syms)`. - syms = vec([symnamedtensor(i, Tuple(axes(tn[i]))) for i in keys(tn)]) - return lazy(Mul(syms)) -end -struct LeftAssociative end -function contraction_order(alg::LeftAssociative, tn) - return prod(i -> symnamedtensor(i, Tuple(axes(tn[i]))), keys(tn)) -end -# Internal implementation shared with the OMEinsumContractionOrders extension. -function _contraction_order(alg, tn) - s = contraction_order(Flat(), tn) - return optimize_evaluation_order(s; alg) -end -function contraction_order(alg::EvaluationOrderAlgorithm, tn) - return _contraction_order(alg, tn) +# A `NormNetwork` has no tensor at a vertex, only the two layers that would form it, so it is +# contracted through `flatten_network`. Supplying a tree alongside it is not meaningful, since the +# tree would have to be keyed by the flattened vertices: pass `flatten_network(nn)` instead, and +# the keys are its own vertices as for any other network. +function contract_network(alg::Exact, nn::NormNetwork) + isnothing(alg.contraction_tree) || throw( + ArgumentError( + "A contraction tree cannot be keyed by the vertices of a `NormNetwork`, whose vertices carry no tensor. Contract `flatten_network(nn)` instead, whose vertices the tree can be built over." + ) + ) + return contract_network(alg, flatten_network(nn)) end diff --git a/src/contraction_tree.jl b/src/contraction_tree.jl new file mode 100644 index 00000000..b69f6d47 --- /dev/null +++ b/src/contraction_tree.jl @@ -0,0 +1,74 @@ +using AbstractTrees: AbstractTrees +using WrappedUnions: @wrapped, unwrap + +""" + ContractionTree{V} + +A fully-resolved binary contraction tree over vertex labels of type `V`. A node is either a +leaf holding a single vertex label, or a branch holding a 2-tuple of child trees; a branch +records that its two children are contracted together, and nesting fixes the pairwise order. + +`ContractionTree` is inert structure, not a lazy contraction: it carries no operation and does +not know how to contract itself. Actually contracting a network maps the leaf labels to tensors +and folds the tree elsewhere; the tree only stores the labels and their nesting. + +Construct a leaf from a label and a branch from two subtrees: + +```julia +ContractionTree(:a) # leaf +ContractionTree(ContractionTree(:a), ContractionTree(:b)) # branch +``` + +Navigate with `isleaf`, integer indexing (`t[1]`, `t[2]`), and the `AbstractTrees` interface +(`children`, `nodevalue`). +""" +@wrapped struct ContractionTree{V} + # A branch's children; equal to `Branch{V}` below, spelled out here because the alias + # references `ContractionTree` and so cannot be defined before this struct. + union::Union{V, NTuple{2, ContractionTree{V}}} +end + +# The children of a branch: a 2-tuple of subtrees. Aliased so the branch member type is named +# everywhere except its definitional site (the struct field above). +const Branch{V} = NTuple{2, ContractionTree{V}} + +# A leaf is built by the default constructor `ContractionTree(label)`. A branch is built from a +# 2-tuple of subtrees; that form is more specific than the default, so `ContractionTree((l, r))` +# builds a branch, not a leaf whose label is a tuple, and `ContractionTree(l, r)` is shorthand. +function ContractionTree(children::Branch{V}) where {V} + return ContractionTree{V}(children) +end +function ContractionTree(left::ContractionTree{V}, right::ContractionTree{V}) where {V} + return ContractionTree{V}((left, right)) +end + +# A branch wraps a 2-tuple of child trees; a leaf wraps its vertex label. Dispatching on the +# branch member type parametrized by `V` (not a bare `Tuple`) keeps this correct even when the +# vertex labels are themselves tuples. +isleaf(t::ContractionTree{V}) where {V} = !(unwrap(t) isa Branch{V}) + +# `nodevalue` is a leaf's vertex label; only a branch has children. +AbstractTrees.nodevalue(t::ContractionTree) = unwrap(t) +AbstractTrees.children(t::ContractionTree) = isleaf(t) ? () : unwrap(t) + +Base.getindex(t::ContractionTree, i::Int) = unwrap(t)[i] + +Base.:(==)(a::ContractionTree, b::ContractionTree) = unwrap(a) == unwrap(b) +Base.hash(t::ContractionTree, h::UInt) = hash(unwrap(t), hash(:ContractionTree, h)) + +function Base.show(io::IO, t::ContractionTree) + if isleaf(t) + show(io, AbstractTrees.nodevalue(t)) + else + print(io, "(") + show(io, t[1]) + print(io, ", ") + show(io, t[2]) + print(io, ")") + end + return nothing +end +function AbstractTrees.printnode(io::IO, t::ContractionTree) + isleaf(t) && show(io, AbstractTrees.nodevalue(t)) + return nothing +end diff --git a/src/normnetwork.jl b/src/normnetwork.jl index 46725507..17d76e86 100644 --- a/src/normnetwork.jl +++ b/src/normnetwork.jl @@ -1,5 +1,5 @@ -using Dictionaries: Dictionary -using ITensorBase: LazyNamedTensor, lazy, rename, setname, similar_operator, uniquename +using Dictionaries: Dictionary, dictionary +using ITensorBase: rename, setname, similar_operator, uniquename using ITensorNetworksNext """ @@ -26,7 +26,11 @@ struct NormNetwork{T, V, I} <: AbstractITensorNetwork{T, V} end end -Base.eltype(::Type{<:NormNetwork{T, V, I}}) where {T, V, I} = LazyNamedTensor{I, T} +function Base.eltype(::Type{<:NormNetwork}) + return error( + "`eltype` of a `NormNetwork` is not defined, since the double-layer tensor at a vertex has no representation of its own. Use `kettensor` and `bratensor` to reach the individual layers." + ) +end function NormNetwork(tn::ITensorNetwork) return NormNetwork(tn, map(uniquename, keys(tn.dimname_vertices))) @@ -46,10 +50,9 @@ NamedGraphs.encoded_graph(nn::NormNetwork) = encoded_graph(nn.ket) # ==================================== DataGraphs.jl ===================================== # function DataGraphs.get_vertex_data(nn::NormNetwork, vertex) - A = kettensor(nn, vertex) - B = conj_bratensor(nn, vertex) - # TODO: implement and use a lazy `conj` via `LazyNamedDimsArrays` here? - return lazy(A) * lazy(conj(B)) + return error( + "Indexing a `NormNetwork` is not defined, since the double-layer tensor at a vertex has no representation of its own. Use `kettensor` and `bratensor` to reach the individual layers." + ) end function DataGraphs.is_vertex_assigned(nn::NormNetwork, vertex) @@ -80,6 +83,29 @@ end bratensor(nn::NormNetwork, vertex) = conj(conj_bratensor(nn, vertex)) +""" + flatten_network(nn::NormNetwork) -> ITensorNetwork + +Expand a norm network into a plain tensor network carrying one vertex per layer, so that vertex +`v` of `nn` becomes the two vertices `(v, :ket)` and `(v, :bra)` and `nv` doubles. A vertex's +two layers stay adjacent in the vertex order, ket first. + +The layers are better contracted as separate operands: the site index they share is invisible +from outside a doubled vertex, so an order built over `nn` must form that vertex before +absorbing anything else. Flattening is explicit rather than automatic because it changes the +vertex set, and a contraction order is only meaningful against the vertices it was built for. +""" +function flatten_network(nn::NormNetwork) + return ITensorNetwork( + dictionary( + Iterators.flatten( + ((v, :ket) => kettensor(nn, v), (v, :bra) => bratensor(nn, v)) + for v in vertices(nn) + ) + ) + ) +end + """ normnetwork(tn::ITensorNetwork, [braname]) -> NormNetwork diff --git a/src/prod_tensors.jl b/src/prod_tensors.jl new file mode 100644 index 00000000..cb3cf79f --- /dev/null +++ b/src/prod_tensors.jl @@ -0,0 +1,81 @@ +using AbstractTrees: nodevalue +using Combinatorics: combinations + +""" + prod_tensors(tensors, tree::ContractionTree) + +Contract `tensors` in the pairwise order given by `tree`, whose leaf labels index into +`tensors`. + +`tensors` is anything indexable by those labels, so a `Vector` pairs with a tree over +positions and a tensor network pairs with a tree over vertices. Separating the order from the +tensors is what lets one tree be reused across many contractions of the same network shape. +""" +function prod_tensors(tensors, tree::ContractionTree) + isleaf(tree) && return tensors[nodevalue(tree)] + return prod_tensors(tensors, tree[1]) * prod_tensors(tensors, tree[2]) +end + +""" + contraction_tree(tensors; alg = Greedy()) + contraction_tree(alg, tensors) + +Find a pairwise contraction order for `tensors`, returned as a [`ContractionTree`](@ref) over +their keys. + +Only the index names and lengths are consulted, never the entries, so an order found once can +be replayed against any network of the same shape. +""" +function contraction_tree end +contraction_tree(tensors; alg = Greedy()) = contraction_tree(alg, tensors) +function contraction_tree(alg, tensors) + return throw(ArgumentError("Contraction order algorithm `$(alg)` not implemented.")) +end + +""" + Greedy + +Repeatedly contract the cheapest available pair, measured as the product of the lengths of all +indices involved. Outer products are taken only once nothing else is left. +""" +struct Greedy end + +""" + LeftAssociative + +Contract in key order, folding from the left. The order `*` would take on its own. +""" +struct LeftAssociative end + +function contraction_tree(::LeftAssociative, tensors) + ks = collect(keys(tensors)) + isempty(ks) && throw(ArgumentError("No tensors to contract.")) + return reduce(ContractionTree, map(ContractionTree, ks)) +end + +function contraction_tree(::Greedy, tensors) + ks = collect(keys(tensors)) + isempty(ks) && throw(ArgumentError("No tensors to contract.")) + trees = map(ContractionTree, ks) + # The axes of each pending operand. An intermediate never materializes, so its axes are + # tracked here rather than read off a tensor. + axs = [collect(axes(tensors[k])) for k in ks] + 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(axs[i], axs[j]) && return typemax(Int) + # Every index on either operand is iterated once, shared ones counted once. + return prod(length, union(axs[i], axs[j]); init = 1) + end + # Shared indices are summed over, so the result carries the symmetric difference. + contracted = symdiff(axs[i1], axs[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]] + axs = [axs[keep]; [contracted]] + end + return only(trees) +end diff --git a/test/Project.toml b/test/Project.toml index 1d3a475e..d761064c 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,4 +1,5 @@ [deps] +AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c" AlgorithmsInterface = "d1e3940c-cd12-4505-8585-b0a4b322527d" Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" DataGraphs = "b5a273c3-7e6c-41f6-98bd-8d7f1525a36a" @@ -23,14 +24,15 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" path = ".." [compat] +AbstractTrees = "0.4.5" AlgorithmsInterface = "0.1" Aqua = "0.8.14" DataGraphs = "0.6" Dictionaries = "0.4.5" GradedArrays = "0.17" Graphs = "1.13.1" -ITensorBase = "0.15.5" -ITensorNetworksNext = "0.10" +ITensorBase = "0.16" +ITensorNetworksNext = "0.11" ITensorPkgSkeleton = "0.3.42" MatrixAlgebraKit = "0.6" NamedGraphs = "0.14" diff --git a/test/test_beliefpropagation.jl b/test/test_beliefpropagation.jl index f838f5d3..189143e2 100644 --- a/test/test_beliefpropagation.jl +++ b/test/test_beliefpropagation.jl @@ -1,15 +1,13 @@ import AlgorithmsInterface as AI -using Base.Broadcast: materialize using DataGraphs: DataGraphs, DataGraph, edge_data, edge_data_type using Dictionaries: Dictionary, dictionary, set! using GradedArrays: U1, gradedrange, isdual using Graphs: AbstractGraph, add_vertex!, dst, edges, has_edge, has_vertex, ne, nv, rem_edge!, src, vertices -using ITensorBase: - Greedy, ITensor, Index, apply, inds, name, noprime, outputnames, prime, state -using ITensorNetworksNext: ITensorNetworksNext, Exact, ITensorNetwork, MessageCache, +using ITensorBase: ITensor, Index, apply, inds, name, noprime, outputnames, prime, state +using ITensorNetworksNext: ITensorNetworksNext, Exact, Greedy, ITensorNetwork, MessageCache, NormNetwork, SimpleMessageUpdate, StopWhenConverged, beliefpropagation, - bethe_free_energy, bethe_free_entropy, bratensor, contract_network, contraction_order, + bethe_free_energy, bethe_free_entropy, bratensor, contract_network, contraction_tree, edge_scalar, edge_scalars, factor_tensors, incoming_messages, insertlink!, kettensor, linkaxes, linkinds, message_environment, messagecache, region_scalar, subgraph, tensornetwork, updated_message, vertex_scalar, vertex_scalars @@ -47,9 +45,9 @@ end struct RecordOperands counts::Vector{Int} end -function ITensorNetworksNext.contraction_order(alg::RecordOperands, tn) +function ITensorNetworksNext.contraction_tree(alg::RecordOperands, tn) push!(alg.counts, length(tn)) - return contraction_order(tn; alg = Greedy()) + return contraction_tree(tn; alg = Greedy()) end @testset "Belief propagation" begin @@ -369,9 +367,8 @@ end nn = NormNetwork(network) v = (2, 2) - # A doubled vertex splits into its two layers, and the split is faithful. - @test length(factor_tensors(nn, v)) == 2 - @test prod(factor_tensors(nn, v)) ≈ materialize(nn[v]) + # A doubled vertex splits into its two layers. + @test factor_tensors(nn, v) == [kettensor(nn, v), bratensor(nn, v)] # A single-layer network's factor is a single operand. @test factor_tensors(network, v) == [network[v]] @@ -380,7 +377,7 @@ end # result matches contracting the doubled vertex as one operand. counts = Int[] algorithm = SimpleMessageUpdate(; - contraction_alg = Exact(; order_alg = RecordOperands(counts)) + contraction_alg = Exact(; tree_alg = RecordOperands(counts)) ) cache = message_environment(one, nn) edge = NamedEdge(v => (2, 3)) @@ -388,7 +385,8 @@ end message = updated_message(algorithm, cache, nn, edge) # `v` has degree 4, so 3 incoming messages plus the ket and bra layers. @test only(counts) == 5 - @test message ≈ contract_network([messages; [nn[v]]]) + @test message ≈ + contract_network([messages; [kettensor(nn, v) * bratensor(nn, v)]]) end end end diff --git a/test/test_contract_network.jl b/test/test_contract_network.jl index 667b7d54..122eefae 100644 --- a/test/test_contract_network.jl +++ b/test/test_contract_network.jl @@ -1,14 +1,13 @@ using Graphs: edges, vertices -using ITensorBase: - Greedy, Index, NamedTensorOperator, inputnames, operator, outputnames, state -using ITensorNetworksNext: Exact, ITensorNetwork, LeftAssociative, contract_network, +using ITensorBase: Index, NamedTensorOperator, inputnames, operator, outputnames, state +using ITensorNetworksNext: Exact, Greedy, ITensorNetwork, LeftAssociative, contract_network, linkinds, siteinds, tensornetwork using NamedGraphs: incident_edges, named_grid using OMEinsumContractionOrders: ExhaustiveSearch, GreedyMethod, TreeSA using Test: @test, @testset @testset "contract_network" begin - orderalg = order_alg -> Exact(; order_alg) + treealg = tree_alg -> Exact(; tree_alg) @testset "Contract Vectors of ITensors" begin i, j, k = Index(2), Index(2), Index(5) @@ -17,11 +16,11 @@ using Test: @test, @testset C = [5.0, 1.0][j] D = [-2.0, 3.0, 4.0, 5.0, 1.0][k] - ABCD_1 = contract_network([A, B, C, D]; alg = orderalg(LeftAssociative())) - ABCD_2 = contract_network([A, B, C, D]; alg = orderalg(Greedy())) - ABCD_3 = contract_network([A, B, C, D]; alg = orderalg(ExhaustiveSearch())) - ABCD_4 = contract_network([A, B, C, D]; alg = orderalg(GreedyMethod())) - ABCD_5 = contract_network([A, B, C, D]; alg = orderalg(TreeSA())) + ABCD_1 = contract_network([A, B, C, D]; alg = treealg(LeftAssociative())) + ABCD_2 = contract_network([A, B, C, D]; alg = treealg(Greedy())) + ABCD_3 = contract_network([A, B, C, D]; alg = treealg(ExhaustiveSearch())) + ABCD_4 = contract_network([A, B, C, D]; alg = treealg(GreedyMethod())) + ABCD_5 = contract_network([A, B, C, D]; alg = treealg(TreeSA())) @test ABCD_1 == ABCD_2 == ABCD_3 @test ABCD_1 ≈ ABCD_4 @test ABCD_1 ≈ ABCD_5 @@ -37,11 +36,11 @@ using Test: @test, @testset return randn(Tuple(is)) end - z1 = contract_network(tn; alg = orderalg(LeftAssociative()))[] - z2 = contract_network(tn; alg = orderalg(Greedy()))[] - z3 = contract_network(tn; alg = orderalg(ExhaustiveSearch()))[] - z4 = contract_network(tn; alg = orderalg(GreedyMethod()))[] - z5 = contract_network(tn; alg = orderalg(TreeSA()))[] + z1 = contract_network(tn; alg = treealg(LeftAssociative()))[] + z2 = contract_network(tn; alg = treealg(Greedy()))[] + z3 = contract_network(tn; alg = treealg(ExhaustiveSearch()))[] + z4 = contract_network(tn; alg = treealg(GreedyMethod()))[] + z5 = contract_network(tn; alg = treealg(TreeSA()))[] @test abs(z1 - z2) / abs(z1) <= 1.0e3 * eps(Float64) @test abs(z1 - z3) / abs(z1) <= 1.0e3 * eps(Float64) diff --git a/test/test_contraction_tree.jl b/test/test_contraction_tree.jl new file mode 100644 index 00000000..542f5016 --- /dev/null +++ b/test/test_contraction_tree.jl @@ -0,0 +1,125 @@ +using AbstractTrees: AbstractTrees, Leaves +using Graphs: edges, vertices +using ITensorBase: Index +using ITensorNetworksNext: ContractionTree, Greedy, LeftAssociative, contraction_tree, + isleaf, prod_tensors, tensornetwork +using NamedGraphs: incident_edges, named_grid +using OMEinsumContractionOrders: GreedyMethod, TreeSA +using Test: @test, @testset + +@testset "ContractionTree" begin + @testset "leaves and branches" begin + l = ContractionTree(:a) + @test isleaf(l) + @test AbstractTrees.nodevalue(l) == :a + @test AbstractTrees.children(l) == () + + # (a, (b, c)) + t = ContractionTree( + ContractionTree(:a), + ContractionTree(ContractionTree(:b), ContractionTree(:c)) + ) + @test !isleaf(t) + @test length(AbstractTrees.children(t)) == 2 + @test t[1] == ContractionTree(:a) + @test isleaf(t[1]) + @test !isleaf(t[2]) + @test AbstractTrees.nodevalue(t[2][1]) == :b + @test AbstractTrees.nodevalue(t[2][2]) == :c + end + + @testset "tuple constructor equals splatted" begin + a, b = ContractionTree(1), ContractionTree(2) + @test ContractionTree((a, b)) == ContractionTree(a, b) + end + + @testset "equality and hashing" begin + t1 = ContractionTree(ContractionTree(:a), ContractionTree(:b)) + t2 = ContractionTree(ContractionTree(:a), ContractionTree(:b)) + t3 = ContractionTree(ContractionTree(:b), ContractionTree(:a)) + @test t1 == t2 + @test hash(t1) == hash(t2) + @test t1 != t3 + # a leaf whose label collides structurally is not equal to a branch + @test ContractionTree(:a) != + ContractionTree(ContractionTree(:a), ContractionTree(:a)) + end + + @testset "arbitrary label types, including Vector-valued leaves" begin + # The ambiguity that nested `Vector` cannot represent: a leaf whose label is itself a + # vector. Leaf/branch is a type distinction, so this is unambiguous. + l = ContractionTree([1, 2]) + @test isleaf(l) + @test AbstractTrees.nodevalue(l) == [1, 2] + + t = ContractionTree(ContractionTree([1, 2]), ContractionTree([3, 4])) + @test !isleaf(t) + @test AbstractTrees.nodevalue(t[1]) == [1, 2] + @test AbstractTrees.nodevalue(t[2]) == [3, 4] + end + + @testset "tuple-valued vertex labels are leaves, not branches" begin + # A `Tuple`-valued label (e.g. grid coordinates) must not be mistaken for a branch: the + # branch arm is `NTuple{2, ContractionTree{V}}`, not any 2-tuple. + l = ContractionTree((1, 2)) + @test isleaf(l) + @test AbstractTrees.nodevalue(l) == (1, 2) + + t = ContractionTree(ContractionTree((1, 2)), ContractionTree((3, 4))) + @test !isleaf(t) + @test isleaf(t[1]) + @test AbstractTrees.nodevalue(t[1]) == (1, 2) + @test AbstractTrees.nodevalue(t[2]) == (3, 4) + end + + @testset "show" begin + t = ContractionTree( + ContractionTree(1), + ContractionTree(ContractionTree(2), ContractionTree(3)) + ) + @test sprint(show, t) == "(1, (2, 3))" + @test sprint(show, ContractionTree(:a)) == ":a" + end + + @testset "`prod_tensors` follows the tree" begin + i, j, k = Index(2), Index(2), Index(5) + A = [1.0 1.0; 0.5 1.0][i, j] + B = [2.0, 1.0][i] + C = [5.0, 1.0][j] + D = [-2.0, 3.0, 4.0, 5.0, 1.0][k] + ts = [A, B, C, D] + + leaves = map(ContractionTree, 1:4) + left = reduce(ContractionTree, leaves) + balanced = ContractionTree( + ContractionTree(leaves[1], leaves[2]), ContractionTree(leaves[3], leaves[4]) + ) + @test prod_tensors(ts, left) ≈ ((A * B) * C) * D + @test prod_tensors(ts, balanced) ≈ (A * B) * (C * D) + end + + @testset "`contraction_tree` over a grid" begin + g = named_grid((2, 3)) + l = Dict(e => Index(2) for e in edges(g)) + l = merge(l, Dict(reverse(e) => l[e] for e in edges(g))) + tn = tensornetwork(vertices(g)) do v + return randn(Tuple(map(e -> l[e], incident_edges(g, v)))) + end + + for alg in (Greedy(), LeftAssociative(), GreedyMethod(), TreeSA()) + tree = contraction_tree(tn; alg) + # Every vertex is contracted exactly once, under its own key. + @test length(collect(Leaves(tree))) == length(vertices(g)) + @test issetequal(map(AbstractTrees.nodevalue, Leaves(tree)), vertices(g)) + @test prod_tensors(tn, tree)[] ≈ prod_tensors(tn, contraction_tree(tn))[] + end + + # A tree depends only on the network's shape, so it replays against a second network + # built over the same graph and indices. + tree = contraction_tree(tn; alg = Greedy()) + tn2 = tensornetwork(vertices(g)) do v + return randn(Tuple(map(e -> l[e], incident_edges(g, v)))) + end + @test prod_tensors(tn2, tree)[] ≈ prod_tensors(tn2, contraction_tree(tn2))[] + end +end diff --git a/test/test_normnetwork.jl b/test/test_normnetwork.jl index 5007be69..9d6cfba1 100644 --- a/test/test_normnetwork.jl +++ b/test/test_normnetwork.jl @@ -1,11 +1,10 @@ using DataGraphs: is_vertex_assigned using Dictionaries: isinsertable, issettable using Graphs: edges, vertices -using ITensorBase: - ITensor, Index, IndexName, LazyITensor, conj, inds, name, setname, uniquename -using ITensorNetworksNext: BraView, ITensorNetwork, KetView, NormNetwork, braname, - bratensor, conj_bratensor, contract_network, indmap, kettensor, normnetwork, - tensornetwork +using ITensorBase: ITensor, Index, IndexName, conj, inds, name, setname, uniquename +using ITensorNetworksNext: BraView, Exact, ITensorNetwork, KetView, NormNetwork, braname, + bratensor, conj_bratensor, contract_network, contraction_tree, flatten_network, indmap, + kettensor, normnetwork, tensornetwork using LinearAlgebra: norm using NamedGraphs: NamedEdge, incident_edges, named_grid, named_path_graph using Test: @test, @test_throws, @testset @@ -37,8 +36,10 @@ end @test issetequal(vertices(nn), vertices(tn)) @test issetequal(edges(nn), edges(tn)) - # `eltype` is the type of the (lazy double-layer) vertex data. - @test eltype(nn) === typeof(nn[1]) + # The double-layer tensor at a vertex has no representation, so neither indexing nor + # `eltype` is defined. + @test_throws ErrorException nn[1] + @test_throws ErrorException eltype(nn) # Vertex data is assigned wherever the underlying network is. @test is_vertex_assigned(nn, 1) @@ -126,6 +127,33 @@ end @test !isinsertable(bv) end + @testset "`flatten_network`" begin + g = named_path_graph(3) + tn, l, s = random_state(Float64, g) + nn = NormNetwork(tn) + flat = flatten_network(nn) + + # One vertex per layer, so the vertex count doubles and each carries its layer's tensor. + @test length(vertices(flat)) == 2 * length(vertices(nn)) + # Each vertex's two layers stay adjacent, in ket-then-bra order. + @test collect(vertices(flat)) == + collect(Iterators.flatten(((v, :ket), (v, :bra)) for v in vertices(nn))) + @test flat[(2, :ket)] === kettensor(nn, 2) + @test flat[(2, :bra)] == bratensor(nn, 2) + + # Contracting the norm network goes through the flattened one. + @test contract_network(nn)[] ≈ contract_network(flat)[] + + # A tree is keyed by the vertices of the network it was built over, so the flattened + # network accepts one and the norm network cannot. + tree = contraction_tree(flat) + @test contract_network(flat; alg = Exact(; contraction_tree = tree))[] ≈ + contract_network(flat)[] + @test_throws ArgumentError contract_network( + nn; alg = Exact(; contraction_tree = tree) + ) + end + @testset "contraction / physics" begin @testset "single normalized tensor contracts to 1" begin s = Index(3) diff --git a/test/test_tensornetwork.jl b/test/test_tensornetwork.jl index d9639475..0beaad01 100644 --- a/test/test_tensornetwork.jl +++ b/test/test_tensornetwork.jl @@ -2,7 +2,7 @@ using DataGraphs: DataGraph, assigned_edge_data, assigned_vertex_data, underlying_graph, vertex_data using Graphs: add_edge!, add_vertex!, dst, edges, edgetype, has_edge, has_vertex, is_directed, ne, nv, rem_edge!, rem_vertex!, src, vertices -using ITensorBase: Index, LazyITensor, inds, operator +using ITensorBase: Index, inds, operator using ITensorNetworksNext: ITensorNetwork, has_ind, linkaxes, linkinds, linknames, operator_support, siteaxes, siteinds, sitenames, tensornetwork using NamedGraphs: convert_vertextype, incident_edges, named_grid, named_path_graph, From 8b4fbb91d32f0171c3713976aabcf0c6c74e65f0 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Mon, 5 Oct 2026 22:16:29 -0400 Subject: [PATCH 2/5] Make `prod_tensors` the single contraction entry point `prod_tensors` now takes either a `ContractionTree` or an algorithm that finds one, so `contract_network` and its `Exact` wrapper are no longer needed. A branch's children get their own `Branch` type, which keeps a leaf distinguishable from a branch even when a vertex label is itself a 2-tuple. Co-Authored-By: Claude Opus 5 (1M context) --- ...etworksNextOMEinsumContractionOrdersExt.jl | 9 +++- src/ITensorNetworksNext.jl | 1 - src/beliefpropagation/beliefpropagation.jl | 13 ++--- src/beliefpropagation/messagecache.jl | 4 +- src/contract_network.jl | 35 ------------ src/contraction_tree.jl | 54 +++++++------------ src/prod_tensors.jl | 52 +++++++++++------- test/test_beliefpropagation.jl | 28 +++++----- test/test_contraction_tree.jl | 27 ++++++---- test/test_itensornetworkgenerators.jl | 6 +-- test/test_normnetwork.jl | 33 ++++++------ ...ntract_network.jl => test_prod_tensors.jl} | 39 +++++++------- 12 files changed, 136 insertions(+), 165 deletions(-) delete mode 100644 src/contract_network.jl rename test/{test_contract_network.jl => test_prod_tensors.jl} (70%) diff --git a/ext/ITensorNetworksNextOMEinsumContractionOrdersExt/ITensorNetworksNextOMEinsumContractionOrdersExt.jl b/ext/ITensorNetworksNextOMEinsumContractionOrdersExt/ITensorNetworksNextOMEinsumContractionOrdersExt.jl index b968975c..2d337a50 100644 --- a/ext/ITensorNetworksNextOMEinsumContractionOrdersExt/ITensorNetworksNextOMEinsumContractionOrdersExt.jl +++ b/ext/ITensorNetworksNextOMEinsumContractionOrdersExt/ITensorNetworksNextOMEinsumContractionOrdersExt.jl @@ -1,7 +1,8 @@ module ITensorNetworksNextOMEinsumContractionOrdersExt using ITensorBase: name -using ITensorNetworksNext: ITensorNetworksNext, ContractionTree +using ITensorNetworksNext: + ITensorNetworksNext, ContractionTree, contraction_tree, prod_tensors using OMEinsumContractionOrders: OMEinsumContractionOrders, CodeOptimizer, EinCode, NestedEinsum, optimize_code @@ -27,4 +28,10 @@ function ITensorNetworksNext.contraction_tree(alg::CodeOptimizer, tensors) return nested_einsum_to_tree(ks, code) end +# `CodeOptimizer` cannot subtype `ContractionTreeAlgorithm`, so the order-taking form of +# `prod_tensors` is extended for it here. +function ITensorNetworksNext.prod_tensors(tensors, alg::CodeOptimizer) + return prod_tensors(tensors, contraction_tree(alg, tensors)) +end + end diff --git a/src/ITensorNetworksNext.jl b/src/ITensorNetworksNext.jl index 7025975d..0298a85f 100644 --- a/src/ITensorNetworksNext.jl +++ b/src/ITensorNetworksNext.jl @@ -18,7 +18,6 @@ include("normnetworkview.jl") include("ITensorNetworkGenerators/ITensorNetworkGenerators.jl") include("contraction_tree.jl") include("prod_tensors.jl") -include("contract_network.jl") include("beliefpropagation/messagecache.jl") include("beliefpropagation/beliefpropagation.jl") diff --git a/src/beliefpropagation/beliefpropagation.jl b/src/beliefpropagation/beliefpropagation.jl index 5a68def3..fa387253 100644 --- a/src/beliefpropagation/beliefpropagation.jl +++ b/src/beliefpropagation/beliefpropagation.jl @@ -232,12 +232,12 @@ function message_update!(cache, factors, edge; alg = nothing, kwargs...) ) end -@kwdef struct SimpleMessageUpdate{ContractionAlg} <: MessageUpdateAlgorithm +@kwdef struct SimpleMessageUpdate{OrderAlg} <: MessageUpdateAlgorithm normalize::Bool = true - contraction_alg::ContractionAlg = Exact() + order_alg::OrderAlg = Greedy() end -# The tensors making up the factor at `vertex`, as separate operands for `contract_network`. A +# The tensors making up the factor at `vertex`, as separate operands to contract. A # `NormNetwork`'s factor is the product of its ket and bra layers. Handing that product over as one # operand would hide the physical index the two layers share, forcing the doubled vertex to be # formed before any message is absorbed (χ^(2 * degree) rather than the χ^(degree + 1) an @@ -251,11 +251,8 @@ end # `edge`. function updated_message(algorithm::SimpleMessageUpdate, cache, factors, edge) messages = collect(incoming_messages(cache, edge)) - # TODO: Remove `factor_tensors` once `contract_network` handles lazy tensors in - # contraction sequences properly. - return contract_network( - [messages; factor_tensors(factors, src(edge))]; alg = algorithm.contraction_alg - ) + tensors = [messages; factor_tensors(factors, src(edge))] + return prod_tensors(tensors, algorithm.order_alg) end # Single-layer network: the message is a plain bond vector, normalized by its entrywise sum. diff --git a/src/beliefpropagation/messagecache.jl b/src/beliefpropagation/messagecache.jl index 650a5858..ec91cd58 100644 --- a/src/beliefpropagation/messagecache.jl +++ b/src/beliefpropagation/messagecache.jl @@ -132,10 +132,8 @@ end function vertex_scalar(factors, messages, vertex; kwargs...) in_messages = incoming_edge_data(messages, [vertex]) - # TODO: Remove `factor_tensors` once `contract_network` handles lazy tensors in - # contraction sequences properly. tensors = [factor_tensors(factors, vertex); collect(in_messages)] - return contract_network(tensors; kwargs...)[] + return prod_tensors(tensors, contraction_tree(tensors; kwargs...))[] end vertex_scalars(factors, messages) = vertex_scalars(factors, messages, keys(factors)) diff --git a/src/contract_network.jl b/src/contract_network.jl deleted file mode 100644 index 981338ba..00000000 --- a/src/contract_network.jl +++ /dev/null @@ -1,35 +0,0 @@ -using Base: @kwdef - -""" - Exact(; contraction_tree = nothing, tree_alg = Greedy()) - -Contract a network exactly, in the pairwise order given by `contraction_tree`. When no tree is -supplied one is found with `tree_alg`. -""" -@kwdef struct Exact{Tree, TreeAlg} - contraction_tree::Tree = nothing - tree_alg::TreeAlg = Greedy() -end - -function contract_network(alg, tn) - return throw(ArgumentError("`contract_network` algorithm `$(alg)` not implemented.")) -end -contract_network(tn; alg = Exact()) = contract_network(alg, tn) - -function contract_network(alg::Exact, tn) - tree = @something alg.contraction_tree contraction_tree(tn; alg = alg.tree_alg) - return prod_tensors(tn, tree) -end - -# A `NormNetwork` has no tensor at a vertex, only the two layers that would form it, so it is -# contracted through `flatten_network`. Supplying a tree alongside it is not meaningful, since the -# tree would have to be keyed by the flattened vertices: pass `flatten_network(nn)` instead, and -# the keys are its own vertices as for any other network. -function contract_network(alg::Exact, nn::NormNetwork) - isnothing(alg.contraction_tree) || throw( - ArgumentError( - "A contraction tree cannot be keyed by the vertices of a `NormNetwork`, whose vertices carry no tensor. Contract `flatten_network(nn)` instead, whose vertices the tree can be built over." - ) - ) - return contract_network(alg, flatten_network(nn)) -end diff --git a/src/contraction_tree.jl b/src/contraction_tree.jl index b69f6d47..6e0c599d 100644 --- a/src/contraction_tree.jl +++ b/src/contraction_tree.jl @@ -2,56 +2,42 @@ using AbstractTrees: AbstractTrees using WrappedUnions: @wrapped, unwrap """ - ContractionTree{V} - -A fully-resolved binary contraction tree over vertex labels of type `V`. A node is either a -leaf holding a single vertex label, or a branch holding a 2-tuple of child trees; a branch -records that its two children are contracted together, and nesting fixes the pairwise order. + Branch(left, right) -`ContractionTree` is inert structure, not a lazy contraction: it carries no operation and does -not know how to contract itself. Actually contracting a network maps the leaf labels to tensors -and folds the tree elsewhere; the tree only stores the labels and their nesting. +The children of a branch node, as held by a [`ContractionTree`](@ref). -Construct a leaf from a label and a branch from two subtrees: +Giving the children a type of their own is what keeps a branch distinguishable from a leaf: a +vertex label can be a 2-tuple, so a bare tuple of children would be ambiguous, but no label is +a `Branch` unless it is deliberately made one. +""" +struct Branch{T} + children::NTuple{2, T} +end +Branch(left, right) = Branch((left, right)) -```julia -ContractionTree(:a) # leaf -ContractionTree(ContractionTree(:a), ContractionTree(:b)) # branch -``` +""" + ContractionTree{V} -Navigate with `isleaf`, integer indexing (`t[1]`, `t[2]`), and the `AbstractTrees` interface -(`children`, `nodevalue`). +A binary tree representing a contraction order over leaf labels of type `V`. A node is either a +leaf holding one label or a [`Branch`](@ref) holding its two child trees. """ @wrapped struct ContractionTree{V} - # A branch's children; equal to `Branch{V}` below, spelled out here because the alias - # references `ContractionTree` and so cannot be defined before this struct. - union::Union{V, NTuple{2, ContractionTree{V}}} + union::Union{V, Branch{ContractionTree{V}}} end -# The children of a branch: a 2-tuple of subtrees. Aliased so the branch member type is named -# everywhere except its definitional site (the struct field above). -const Branch{V} = NTuple{2, ContractionTree{V}} - -# A leaf is built by the default constructor `ContractionTree(label)`. A branch is built from a -# 2-tuple of subtrees; that form is more specific than the default, so `ContractionTree((l, r))` -# builds a branch, not a leaf whose label is a tuple, and `ContractionTree(l, r)` is shorthand. -function ContractionTree(children::Branch{V}) where {V} +function ContractionTree(children::Branch{ContractionTree{V}}) where {V} return ContractionTree{V}(children) end function ContractionTree(left::ContractionTree{V}, right::ContractionTree{V}) where {V} - return ContractionTree{V}((left, right)) + return ContractionTree{V}(Branch(left, right)) end -# A branch wraps a 2-tuple of child trees; a leaf wraps its vertex label. Dispatching on the -# branch member type parametrized by `V` (not a bare `Tuple`) keeps this correct even when the -# vertex labels are themselves tuples. -isleaf(t::ContractionTree{V}) where {V} = !(unwrap(t) isa Branch{V}) +isleaf(t::ContractionTree) = !(unwrap(t) isa Branch) -# `nodevalue` is a leaf's vertex label; only a branch has children. AbstractTrees.nodevalue(t::ContractionTree) = unwrap(t) -AbstractTrees.children(t::ContractionTree) = isleaf(t) ? () : unwrap(t) +AbstractTrees.children(t::ContractionTree) = isleaf(t) ? () : unwrap(t).children -Base.getindex(t::ContractionTree, i::Int) = unwrap(t)[i] +Base.getindex(t::ContractionTree, i::Int) = unwrap(t).children[i] Base.:(==)(a::ContractionTree, b::ContractionTree) = unwrap(a) == unwrap(b) Base.hash(t::ContractionTree, h::UInt) = hash(unwrap(t), hash(:ContractionTree, h)) diff --git a/src/prod_tensors.jl b/src/prod_tensors.jl index cb3cf79f..3f731968 100644 --- a/src/prod_tensors.jl +++ b/src/prod_tensors.jl @@ -1,20 +1,49 @@ using AbstractTrees: nodevalue using Combinatorics: combinations +""" + ContractionTreeAlgorithm + +Supertype of the strategies for finding a contraction order. See [`contraction_tree`](@ref). +""" +abstract type ContractionTreeAlgorithm <: AbstractAlgorithm end + """ prod_tensors(tensors, tree::ContractionTree) + prod_tensors(tensors, alg) + prod_tensors(tensors) Contract `tensors` in the pairwise order given by `tree`, whose leaf labels index into -`tensors`. +`tensors`. Given an order algorithm instead, find a tree with it first. + +`tensors` is anything indexable by those labels, so a `Vector` pairs with a tree over positions +and a tensor network pairs with a tree over vertices. Separating the order from the tensors is +what lets one tree be reused across many contractions of the same network shape. -`tensors` is anything indexable by those labels, so a `Vector` pairs with a tree over -positions and a tensor network pairs with a tree over vertices. Separating the order from the -tensors is what lets one tree be reused across many contractions of the same network shape. +With neither, the tensors are folded from the left, the order `*` would take on its own. """ +function prod_tensors end +prod_tensors(tensors) = prod_tensors(tensors, left_associative_tree(keys(tensors))) function prod_tensors(tensors, tree::ContractionTree) isleaf(tree) && return tensors[nodevalue(tree)] return prod_tensors(tensors, tree[1]) * prod_tensors(tensors, tree[2]) end +function prod_tensors(tensors, alg::ContractionTreeAlgorithm) + return prod_tensors(tensors, contraction_tree(alg, tensors)) +end + +""" + left_associative_tree(labels) + +The tree that folds `labels` from the left, so `(a, b, c)` gives `((a, b), c)`. + +Unlike [`contraction_tree`](@ref) this consults nothing but the labels themselves, since the +order is fixed rather than searched for. +""" +function left_associative_tree(labels) + isempty(labels) && throw(ArgumentError("No tensors to contract.")) + return reduce(ContractionTree, map(ContractionTree, labels)) +end """ contraction_tree(tensors; alg = Greedy()) @@ -38,20 +67,7 @@ end Repeatedly contract the cheapest available pair, measured as the product of the lengths of all indices involved. Outer products are taken only once nothing else is left. """ -struct Greedy end - -""" - LeftAssociative - -Contract in key order, folding from the left. The order `*` would take on its own. -""" -struct LeftAssociative end - -function contraction_tree(::LeftAssociative, tensors) - ks = collect(keys(tensors)) - isempty(ks) && throw(ArgumentError("No tensors to contract.")) - return reduce(ContractionTree, map(ContractionTree, ks)) -end +struct Greedy <: ContractionTreeAlgorithm end function contraction_tree(::Greedy, tensors) ks = collect(keys(tensors)) diff --git a/test/test_beliefpropagation.jl b/test/test_beliefpropagation.jl index 189143e2..1591e26f 100644 --- a/test/test_beliefpropagation.jl +++ b/test/test_beliefpropagation.jl @@ -5,12 +5,12 @@ using GradedArrays: U1, gradedrange, isdual using Graphs: AbstractGraph, add_vertex!, dst, edges, has_edge, has_vertex, ne, nv, rem_edge!, src, vertices using ITensorBase: ITensor, Index, apply, inds, name, noprime, outputnames, prime, state -using ITensorNetworksNext: ITensorNetworksNext, Exact, Greedy, ITensorNetwork, MessageCache, - NormNetwork, SimpleMessageUpdate, StopWhenConverged, beliefpropagation, - bethe_free_energy, bethe_free_entropy, bratensor, contract_network, contraction_tree, +using ITensorNetworksNext: ITensorNetworksNext, ContractionTreeAlgorithm, Greedy, + ITensorNetwork, MessageCache, NormNetwork, SimpleMessageUpdate, StopWhenConverged, + beliefpropagation, bethe_free_energy, bethe_free_entropy, bratensor, contraction_tree, edge_scalar, edge_scalars, factor_tensors, incoming_messages, insertlink!, kettensor, - linkaxes, linkinds, message_environment, messagecache, region_scalar, subgraph, - tensornetwork, updated_message, vertex_scalar, vertex_scalars + linkaxes, linkinds, message_environment, messagecache, prod_tensors, region_scalar, + subgraph, tensornetwork, updated_message, vertex_scalar, vertex_scalars using LinearAlgebra: LinearAlgebra, norm, tr using NamedGraphs: NamedEdge, all_edges, incident_edges, named_comb_tree, named_grid, named_path_graph, vertextype @@ -41,8 +41,8 @@ function spin_ice_tensornetwork(g) return ITensorNetwork(ts) end -# Records how many operands each `contract_network` call is given, then orders them greedily. -struct RecordOperands +# Records how many operands each contraction is given, then orders them greedily. +struct RecordOperands <: ContractionTreeAlgorithm counts::Vector{Int} end function ITensorNetworksNext.contraction_tree(alg::RecordOperands, tn) @@ -343,7 +343,7 @@ end for (edge, rest) in ((1 => 2, 2:4), (4 => 3, 1:3)) layers = [[kettensor(nn, v) for v in rest]; [bratensor(nn, v) for v in rest]] - z_rest = contract_network([state(ones[edge]); layers])[] + z_rest = prod_tensors([state(ones[edge]); layers], Greedy())[] @test z_rest ≈ norm(prod([network[v] for v in rest]))^2 rtol = eps(real(T))^(1 / 3) end @@ -372,13 +372,11 @@ end # A single-layer network's factor is a single operand. @test factor_tensors(network, v) == [network[v]] - # The message update passes the layers to `contract_network` as separate operands, so - # the contraction order can interleave the incoming messages between them, and the - # result matches contracting the doubled vertex as one operand. + # The message update passes the layers as separate operands, so the contraction + # order can interleave the incoming messages between them, and the result matches + # contracting the doubled vertex as one operand. counts = Int[] - algorithm = SimpleMessageUpdate(; - contraction_alg = Exact(; tree_alg = RecordOperands(counts)) - ) + algorithm = SimpleMessageUpdate(; order_alg = RecordOperands(counts)) cache = message_environment(one, nn) edge = NamedEdge(v => (2, 3)) messages = collect(incoming_messages(cache, edge)) @@ -386,7 +384,7 @@ end # `v` has degree 4, so 3 incoming messages plus the ket and bra layers. @test only(counts) == 5 @test message ≈ - contract_network([messages; [kettensor(nn, v) * bratensor(nn, v)]]) + prod_tensors([messages; [kettensor(nn, v) * bratensor(nn, v)]], Greedy()) end end end diff --git a/test/test_contraction_tree.jl b/test/test_contraction_tree.jl index 542f5016..54538553 100644 --- a/test/test_contraction_tree.jl +++ b/test/test_contraction_tree.jl @@ -1,8 +1,8 @@ using AbstractTrees: AbstractTrees, Leaves using Graphs: edges, vertices using ITensorBase: Index -using ITensorNetworksNext: ContractionTree, Greedy, LeftAssociative, contraction_tree, - isleaf, prod_tensors, tensornetwork +using ITensorNetworksNext: Branch, ContractionTree, Greedy, contraction_tree, isleaf, + left_associative_tree, prod_tensors, tensornetwork using NamedGraphs: incident_edges, named_grid using OMEinsumContractionOrders: GreedyMethod, TreeSA using Test: @test, @testset @@ -28,9 +28,15 @@ using Test: @test, @testset @test AbstractTrees.nodevalue(t[2][2]) == :c end - @testset "tuple constructor equals splatted" begin + @testset "a branch comes from `Branch`, never from a bare tuple" begin a, b = ContractionTree(1), ContractionTree(2) - @test ContractionTree((a, b)) == ContractionTree(a, b) + @test ContractionTree(Branch(a, b)) == ContractionTree(a, b) + @test !isleaf(ContractionTree(Branch(a, b))) + + # A bare 2-tuple is a label, even when its entries are themselves trees. This is what + # keeps a tuple-valued vertex label from being read as children. + @test isleaf(ContractionTree((a, b))) + @test AbstractTrees.nodevalue(ContractionTree((a, b))) == (a, b) end @testset "equality and hashing" begin @@ -59,8 +65,8 @@ using Test: @test, @testset end @testset "tuple-valued vertex labels are leaves, not branches" begin - # A `Tuple`-valued label (e.g. grid coordinates) must not be mistaken for a branch: the - # branch arm is `NTuple{2, ContractionTree{V}}`, not any 2-tuple. + # A `Tuple`-valued label (e.g. grid coordinates) must not be mistaken for a branch. + # Leaves and branches are separate types, so a label can never be read as children. l = ContractionTree((1, 2)) @test isleaf(l) @test AbstractTrees.nodevalue(l) == (1, 2) @@ -90,7 +96,8 @@ using Test: @test, @testset ts = [A, B, C, D] leaves = map(ContractionTree, 1:4) - left = reduce(ContractionTree, leaves) + left = left_associative_tree(1:4) + @test left == reduce(ContractionTree, leaves) balanced = ContractionTree( ContractionTree(leaves[1], leaves[2]), ContractionTree(leaves[3], leaves[4]) ) @@ -106,12 +113,12 @@ using Test: @test, @testset return randn(Tuple(map(e -> l[e], incident_edges(g, v)))) end - for alg in (Greedy(), LeftAssociative(), GreedyMethod(), TreeSA()) + for alg in (Greedy(), GreedyMethod(), TreeSA()) tree = contraction_tree(tn; alg) # Every vertex is contracted exactly once, under its own key. @test length(collect(Leaves(tree))) == length(vertices(g)) @test issetequal(map(AbstractTrees.nodevalue, Leaves(tree)), vertices(g)) - @test prod_tensors(tn, tree)[] ≈ prod_tensors(tn, contraction_tree(tn))[] + @test prod_tensors(tn, tree)[] ≈ prod_tensors(tn, Greedy())[] end # A tree depends only on the network's shape, so it replays against a second network @@ -120,6 +127,6 @@ using Test: @test, @testset tn2 = tensornetwork(vertices(g)) do v return randn(Tuple(map(e -> l[e], incident_edges(g, v)))) end - @test prod_tensors(tn2, tree)[] ≈ prod_tensors(tn2, contraction_tree(tn2))[] + @test prod_tensors(tn2, tree)[] ≈ prod_tensors(tn2, Greedy())[] end end diff --git a/test/test_itensornetworkgenerators.jl b/test/test_itensornetworkgenerators.jl index be154eb1..0dc9f4b7 100644 --- a/test/test_itensornetworkgenerators.jl +++ b/test/test_itensornetworkgenerators.jl @@ -2,7 +2,7 @@ using DataGraphs: DataGraph using Graphs: edges, ne, nv, vertices using ITensorBase: Index, inds using ITensorNetworksNext.ITensorNetworkGenerators: delta, delta_network, ising_network -using ITensorNetworksNext: contract_network +using ITensorNetworksNext: Greedy, prod_tensors using NamedGraphs: arranged_edges, incident_edges, named_grid using Test: @test, @testset @@ -46,7 +46,7 @@ using Test: @test, @testset @test issetequal(is, inds(tn[v])) @test tn[v] ≠ delta(Float64, Tuple(is)) end - z = contract_network(tn)[] + z = prod_tensors(tn, Greedy())[] f = -log(z) / (β * nv(g)) f_analytic = TestUtils.f_1d_ising(β, 4; periodic) @test f ≈ f_analytic @@ -67,7 +67,7 @@ using Test: @test, @testset @test issetequal(is, inds(tn[v])) @test tn[v] ≠ delta(Float64, Tuple(is)) end - z = contract_network(tn)[] + z = prod_tensors(tn, Greedy())[] f = -log(z) / (β * nv(g)) f_inf = TestUtils.f_2d_ising(β) @test f ≈ f_inf rtol = 1.0e-1 diff --git a/test/test_normnetwork.jl b/test/test_normnetwork.jl index 9d6cfba1..4528b220 100644 --- a/test/test_normnetwork.jl +++ b/test/test_normnetwork.jl @@ -2,9 +2,9 @@ using DataGraphs: is_vertex_assigned using Dictionaries: isinsertable, issettable using Graphs: edges, vertices using ITensorBase: ITensor, Index, IndexName, conj, inds, name, setname, uniquename -using ITensorNetworksNext: BraView, Exact, ITensorNetwork, KetView, NormNetwork, braname, - bratensor, conj_bratensor, contract_network, contraction_tree, flatten_network, indmap, - kettensor, normnetwork, tensornetwork +using ITensorNetworksNext: BraView, Greedy, ITensorNetwork, KetView, NormNetwork, braname, + bratensor, conj_bratensor, contraction_tree, flatten_network, indmap, kettensor, + normnetwork, prod_tensors, tensornetwork using LinearAlgebra: norm using NamedGraphs: NamedEdge, incident_edges, named_grid, named_path_graph using Test: @test, @test_throws, @testset @@ -22,6 +22,10 @@ function random_state(::Type{T}, g; d = 2, χ = 2) where {T} return tn, l, s end +# A norm network carries no tensor at a vertex, so contracting it goes through the flattened +# network, whose two layers are separate operands. +contract_norm(nn) = prod_tensors(flatten_network(nn), Greedy())[] + @testset "`NormNetwork`" begin @testset "Basics" begin g = named_path_graph(3) @@ -141,17 +145,13 @@ end @test flat[(2, :ket)] === kettensor(nn, 2) @test flat[(2, :bra)] == bratensor(nn, 2) - # Contracting the norm network goes through the flattened one. - @test contract_network(nn)[] ≈ contract_network(flat)[] + # A norm network carries no tensor at a vertex, so it has to be flattened first. + @test_throws ErrorException prod_tensors(nn, Greedy()) - # A tree is keyed by the vertices of the network it was built over, so the flattened - # network accepts one and the norm network cannot. + # A tree is keyed by the vertices of the network it was built over, so one built over + # the flattened network replays against it. tree = contraction_tree(flat) - @test contract_network(flat; alg = Exact(; contraction_tree = tree))[] ≈ - contract_network(flat)[] - @test_throws ArgumentError contract_network( - nn; alg = Exact(; contraction_tree = tree) - ) + @test prod_tensors(flat, tree)[] ≈ prod_tensors(flat, Greedy())[] end @testset "contraction / physics" begin @@ -163,7 +163,7 @@ end nn = NormNetwork(tn) # ⟨ψ|ψ⟩ for a single normalized site tensor is 1. - @test contract_network(nn)[] ≈ 1 + @test contract_norm(nn) ≈ 1 end @testset "$T" for T in (Float64, ComplexF64) @@ -171,19 +171,18 @@ end tn, l, s = random_state(T, g) # The norm network contracts to ⟨tn|tn⟩ = ‖prod(tn)‖², a real nonnegative number. - z = contract_network(NormNetwork(tn))[] + z = contract_norm(NormNetwork(tn)) @test z ≈ norm(prod(tn))^2 @test imag(z) ≈ 0 atol = 1.0e-12 * abs(z) @test real(z) > 0 # Rescaling a single tensor by 1/√z normalizes the state, so ⟨tn|tn⟩ = 1. tn[first(vertices(tn))] = tn[first(vertices(tn))] / sqrt(real(z)) - @test contract_network(NormNetwork(tn))[] ≈ 1 + @test contract_norm(NormNetwork(tn)) ≈ 1 # The contracted norm does not depend on the chosen bra-layer name map. custom = map(uniquename, keys(tn.dimname_vertices)) - @test contract_network(normnetwork(tn, custom))[] ≈ - contract_network(NormNetwork(tn))[] + @test contract_norm(normnetwork(tn, custom)) ≈ contract_norm(NormNetwork(tn)) end end end diff --git a/test/test_contract_network.jl b/test/test_prod_tensors.jl similarity index 70% rename from test/test_contract_network.jl rename to test/test_prod_tensors.jl index 122eefae..b31d12bd 100644 --- a/test/test_contract_network.jl +++ b/test/test_prod_tensors.jl @@ -1,14 +1,12 @@ using Graphs: edges, vertices using ITensorBase: Index, NamedTensorOperator, inputnames, operator, outputnames, state -using ITensorNetworksNext: Exact, Greedy, ITensorNetwork, LeftAssociative, contract_network, - linkinds, siteinds, tensornetwork +using ITensorNetworksNext: + Greedy, ITensorNetwork, linkinds, prod_tensors, siteinds, tensornetwork using NamedGraphs: incident_edges, named_grid using OMEinsumContractionOrders: ExhaustiveSearch, GreedyMethod, TreeSA using Test: @test, @testset -@testset "contract_network" begin - treealg = tree_alg -> Exact(; tree_alg) - +@testset "prod_tensors" begin @testset "Contract Vectors of ITensors" begin i, j, k = Index(2), Index(2), Index(5) A = [1.0 1.0; 0.5 1.0][i, j] @@ -16,11 +14,12 @@ using Test: @test, @testset C = [5.0, 1.0][j] D = [-2.0, 3.0, 4.0, 5.0, 1.0][k] - ABCD_1 = contract_network([A, B, C, D]; alg = treealg(LeftAssociative())) - ABCD_2 = contract_network([A, B, C, D]; alg = treealg(Greedy())) - ABCD_3 = contract_network([A, B, C, D]; alg = treealg(ExhaustiveSearch())) - ABCD_4 = contract_network([A, B, C, D]; alg = treealg(GreedyMethod())) - ABCD_5 = contract_network([A, B, C, D]; alg = treealg(TreeSA())) + ts = [A, B, C, D] + ABCD_1 = prod_tensors(ts) + ABCD_2 = prod_tensors(ts, Greedy()) + ABCD_3 = prod_tensors(ts, ExhaustiveSearch()) + ABCD_4 = prod_tensors(ts, GreedyMethod()) + ABCD_5 = prod_tensors(ts, TreeSA()) @test ABCD_1 == ABCD_2 == ABCD_3 @test ABCD_1 ≈ ABCD_4 @test ABCD_1 ≈ ABCD_5 @@ -36,11 +35,11 @@ using Test: @test, @testset return randn(Tuple(is)) end - z1 = contract_network(tn; alg = treealg(LeftAssociative()))[] - z2 = contract_network(tn; alg = treealg(Greedy()))[] - z3 = contract_network(tn; alg = treealg(ExhaustiveSearch()))[] - z4 = contract_network(tn; alg = treealg(GreedyMethod()))[] - z5 = contract_network(tn; alg = treealg(TreeSA()))[] + z1 = prod_tensors(tn)[] + z2 = prod_tensors(tn, Greedy())[] + z3 = prod_tensors(tn, ExhaustiveSearch())[] + z4 = prod_tensors(tn, GreedyMethod())[] + z5 = prod_tensors(tn, TreeSA())[] @test abs(z1 - z2) / abs(z1) <= 1.0e3 * eps(Float64) @test abs(z1 - z3) / abs(z1) <= 1.0e3 * eps(Float64) @@ -58,7 +57,7 @@ using Test: @test, @testset # A network mixing an operator with plain tensors previously threw a `convert` # `MethodError`; it now contracts, stays an operator, and matches the binary product. t = randn(2, 2)[j, k] - r = contract_network([o, t]) + r = prod_tensors([o, t]) @test r isa NamedTensorOperator @test state(r) ≈ state(o * t) @test outputnames(r) == outputnames(o * t) @@ -66,18 +65,18 @@ using Test: @test, @testset # An all-operator network is likewise preserved. o2 = operator(randn(2, 2), (j,), (k,)) - r2 = contract_network([o, o2]) + r2 = prod_tensors([o, o2]) @test r2 isa NamedTensorOperator @test state(r2) ≈ state(o * o2) # A fully-contracted operator network reads out as a scalar via `[]`. f = randn(2, 2)[i, j] - @test contract_network([o, f])[] ≈ (o * f)[] + @test prod_tensors([o, f])[] ≈ (o * f)[] # An all-plain network is unaffected: it is not promoted to an operator. a = randn(2, 2)[i, j] b = randn(2, 2)[j, k] - @test !(contract_network([a, b]) isa NamedTensorOperator) + @test !(prod_tensors([a, b]) isa NamedTensorOperator) # Pairing is order-independent: a branching network with a surviving output/input pair # matches the binary product under any fold order (greedy contraction included). @@ -85,7 +84,7 @@ using Test: @test, @testset op = operator(randn(2, 2, 2, 2), (ip, mp), (i, m)) u = randn(2, 2)[mp, x] w = randn(2, 2)[m, x] - rb = contract_network([op, u, w]) + rb = prod_tensors([op, u, w]) @test outputnames(rb) == outputnames((op * u) * w) == outputnames(op * (u * w)) @test inputnames(rb) == inputnames((op * u) * w) == inputnames(op * (u * w)) end From 1e99951efce28943b1dd81811eca49dcdff050eb Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Tue, 6 Oct 2026 12:28:23 -0400 Subject: [PATCH 3/5] Trim the contraction-order layer `ContractionTree` wrapped its two-arm union with `WrappedUnions` but only ever called `unwrap`, so a plain struct does the same job without the dependency. The greedy search now reads `inds`, the name this stack uses for a tensor's indices, rather than going through `axes`. Co-Authored-By: Claude Opus 5 (1M context) --- Project.toml | 2 -- src/contraction_tree.jl | 17 ++++++++--------- src/normnetwork.jl | 5 ----- src/prod_tensors.jl | 30 +++++++++++------------------- 4 files changed, 19 insertions(+), 35 deletions(-) diff --git a/Project.toml b/Project.toml index 500e476c..ace39095 100644 --- a/Project.toml +++ b/Project.toml @@ -23,7 +23,6 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" SimpleTraits = "699a6c99-e7fa-54fc-8d76-47d257e15c1d" SplitApplyCombine = "03a91e81-4c3e-53e1-a0a4-9c0c8f19dd66" TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a" -WrappedUnions = "325db55a-9c6c-5b90-b1a2-ec87e7a38c44" [weakdeps] OMEinsumContractionOrders = "6f22d1fd-8eed-4bb7-9776-e7d684900715" @@ -53,5 +52,4 @@ Random = "1.10" SimpleTraits = "0.9.5" SplitApplyCombine = "1.2.3" TensorAlgebra = "0.23" -WrappedUnions = "0.3" julia = "1.10" diff --git a/src/contraction_tree.jl b/src/contraction_tree.jl index 6e0c599d..0fe6d9e1 100644 --- a/src/contraction_tree.jl +++ b/src/contraction_tree.jl @@ -1,5 +1,4 @@ using AbstractTrees: AbstractTrees -using WrappedUnions: @wrapped, unwrap """ Branch(left, right) @@ -21,8 +20,8 @@ Branch(left, right) = Branch((left, right)) A binary tree representing a contraction order over leaf labels of type `V`. A node is either a leaf holding one label or a [`Branch`](@ref) holding its two child trees. """ -@wrapped struct ContractionTree{V} - union::Union{V, Branch{ContractionTree{V}}} +struct ContractionTree{V} + node::Union{V, Branch{ContractionTree{V}}} end function ContractionTree(children::Branch{ContractionTree{V}}) where {V} @@ -32,15 +31,15 @@ function ContractionTree(left::ContractionTree{V}, right::ContractionTree{V}) wh return ContractionTree{V}(Branch(left, right)) end -isleaf(t::ContractionTree) = !(unwrap(t) isa Branch) +isleaf(t::ContractionTree) = !(t.node isa Branch) -AbstractTrees.nodevalue(t::ContractionTree) = unwrap(t) -AbstractTrees.children(t::ContractionTree) = isleaf(t) ? () : unwrap(t).children +AbstractTrees.nodevalue(t::ContractionTree) = t.node +AbstractTrees.children(t::ContractionTree) = isleaf(t) ? () : t.node.children -Base.getindex(t::ContractionTree, i::Int) = unwrap(t).children[i] +Base.getindex(t::ContractionTree, i::Int) = t.node.children[i] -Base.:(==)(a::ContractionTree, b::ContractionTree) = unwrap(a) == unwrap(b) -Base.hash(t::ContractionTree, h::UInt) = hash(unwrap(t), hash(:ContractionTree, h)) +Base.:(==)(a::ContractionTree, b::ContractionTree) = a.node == b.node +Base.hash(t::ContractionTree, h::UInt) = hash(t.node, hash(:ContractionTree, h)) function Base.show(io::IO, t::ContractionTree) if isleaf(t) diff --git a/src/normnetwork.jl b/src/normnetwork.jl index 17d76e86..d735b22f 100644 --- a/src/normnetwork.jl +++ b/src/normnetwork.jl @@ -89,11 +89,6 @@ bratensor(nn::NormNetwork, vertex) = conj(conj_bratensor(nn, vertex)) Expand a norm network into a plain tensor network carrying one vertex per layer, so that vertex `v` of `nn` becomes the two vertices `(v, :ket)` and `(v, :bra)` and `nv` doubles. A vertex's two layers stay adjacent in the vertex order, ket first. - -The layers are better contracted as separate operands: the site index they share is invisible -from outside a doubled vertex, so an order built over `nn` must form that vertex before -absorbing anything else. Flattening is explicit rather than automatic because it changes the -vertex set, and a contraction order is only meaningful against the vertices it was built for. """ function flatten_network(nn::NormNetwork) return ITensorNetwork( diff --git a/src/prod_tensors.jl b/src/prod_tensors.jl index 3f731968..a9d8a879 100644 --- a/src/prod_tensors.jl +++ b/src/prod_tensors.jl @@ -1,5 +1,6 @@ using AbstractTrees: nodevalue using Combinatorics: combinations +using ITensorBase: inds """ ContractionTreeAlgorithm @@ -14,13 +15,11 @@ abstract type ContractionTreeAlgorithm <: AbstractAlgorithm end prod_tensors(tensors) Contract `tensors` in the pairwise order given by `tree`, whose leaf labels index into -`tensors`. Given an order algorithm instead, find a tree with it first. +`tensors`. Given an order algorithm instead, find a tree with it first. Given neither, contract +in key order, folding from the left. `tensors` is anything indexable by those labels, so a `Vector` pairs with a tree over positions -and a tensor network pairs with a tree over vertices. Separating the order from the tensors is -what lets one tree be reused across many contractions of the same network shape. - -With neither, the tensors are folded from the left, the order `*` would take on its own. +and a tensor network pairs with a tree over vertices. """ function prod_tensors end prod_tensors(tensors) = prod_tensors(tensors, left_associative_tree(keys(tensors))) @@ -32,14 +31,7 @@ function prod_tensors(tensors, alg::ContractionTreeAlgorithm) return prod_tensors(tensors, contraction_tree(alg, tensors)) end -""" - left_associative_tree(labels) - -The tree that folds `labels` from the left, so `(a, b, c)` gives `((a, b), c)`. - -Unlike [`contraction_tree`](@ref) this consults nothing but the labels themselves, since the -order is fixed rather than searched for. -""" +# The tree that folds `labels` from the left, so `(a, b, c)` gives `((a, b), c)`. function left_associative_tree(labels) isempty(labels) && throw(ArgumentError("No tensors to contract.")) return reduce(ContractionTree, map(ContractionTree, labels)) @@ -73,25 +65,25 @@ function contraction_tree(::Greedy, tensors) ks = collect(keys(tensors)) isempty(ks) && throw(ArgumentError("No tensors to contract.")) trees = map(ContractionTree, ks) - # The axes of each pending operand. An intermediate never materializes, so its axes are + # The inds of each pending operand. An intermediate never materializes, so its inds are # tracked here rather than read off a tensor. - axs = [collect(axes(tensors[k])) for k in ks] + is = [inds(tensors[k]) for k in ks] 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(axs[i], axs[j]) && return typemax(Int) + isdisjoint(is[i], is[j]) && return typemax(Int) # Every index on either operand is iterated once, shared ones counted once. - return prod(length, union(axs[i], axs[j]); init = 1) + return prod(length, union(is[i], is[j]); init = 1) end # Shared indices are summed over, so the result carries the symmetric difference. - contracted = symdiff(axs[i1], axs[i2]) + contracted = symdiff(is[i1], 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]] - axs = [axs[keep]; [contracted]] + is = [is[keep]; [contracted]] end return only(trees) end From 0da075032fb2baff691d158b521eae3b8dff5578 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Tue, 6 Oct 2026 20:43:00 -0400 Subject: [PATCH 4/5] Resolve ITensorBase from the registry ITensorBase 0.16.0 is registered, so the branch pin that let CI resolve the unreleased dependency is no longer needed. Co-Authored-By: Claude Opus 5 (1M context) --- Project.toml | 4 ---- 1 file changed, 4 deletions(-) diff --git a/Project.toml b/Project.toml index ace39095..5b863052 100644 --- a/Project.toml +++ b/Project.toml @@ -27,10 +27,6 @@ TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a" [weakdeps] OMEinsumContractionOrders = "6f22d1fd-8eed-4bb7-9776-e7d684900715" -[sources.ITensorBase] -rev = "mf/remove-lazy" -url = "https://github.com/ITensor/ITensorBase.jl" - [extensions] ITensorNetworksNextOMEinsumContractionOrdersExt = "OMEinsumContractionOrders" From 0ff41908840f4c2d50f4d5233c28bcd0f2e821fb Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Tue, 6 Oct 2026 21:00:14 -0400 Subject: [PATCH 5/5] Use `inds` in the contraction-order extension Matches the greedy search, which already reads `inds` rather than going through `axes`. Co-Authored-By: Claude Opus 5 (1M context) --- .../ITensorNetworksNextOMEinsumContractionOrdersExt.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/ext/ITensorNetworksNextOMEinsumContractionOrdersExt/ITensorNetworksNextOMEinsumContractionOrdersExt.jl b/ext/ITensorNetworksNextOMEinsumContractionOrdersExt/ITensorNetworksNextOMEinsumContractionOrdersExt.jl index 2d337a50..ca80cf7c 100644 --- a/ext/ITensorNetworksNextOMEinsumContractionOrdersExt/ITensorNetworksNextOMEinsumContractionOrdersExt.jl +++ b/ext/ITensorNetworksNextOMEinsumContractionOrdersExt/ITensorNetworksNextOMEinsumContractionOrdersExt.jl @@ -1,6 +1,6 @@ module ITensorNetworksNextOMEinsumContractionOrdersExt -using ITensorBase: name +using ITensorBase: inds, name using ITensorNetworksNext: ITensorNetworksNext, ContractionTree, contraction_tree, prod_tensors using OMEinsumContractionOrders: @@ -18,10 +18,10 @@ end # `TreeSA`, `KaHyParBipartite`, ...) by forwarding to `optimize_code`. function ITensorNetworksNext.contraction_tree(alg::CodeOptimizer, tensors) ks = collect(keys(tensors)) - ixs = [map(name, collect(axes(tensors[k]))) for k in ks] + ixs = [map(name, inds(tensors[k])) for k in ks] all_inds = reduce(vcat, ixs) labels = unique(all_inds) - size_dict = Dict(name(ax) => length(ax) for k in ks for ax in axes(tensors[k])) + size_dict = Dict(name(i) => length(i) for k in ks for i in inds(tensors[k])) # Open indices (appearing on a single tensor) are the output of the network. iy = filter(i -> count(==(i), all_inds) == 1, labels) code = optimize_code(EinCode(ixs, iy), size_dict, alg)