From 73c0cae70a0d400a2eb7d3e15aa4efda99bd8b39 Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Wed, 29 Jul 2026 09:48:18 -0400 Subject: [PATCH] Flatten lazy product operands into leaves for contraction ordering MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `contract_network` lowered each operand to a single symbol carrying its outer axes, so a lazy product operand stayed opaque to the contraction-order optimizer. A `NormNetwork`'s doubled vertex is such a product, `lazy(ket) * lazy(conj(bra))`, and collapsing it hid the physical index shared between the ket and bra. The optimizer could not interleave the other operands between the two layers, so it formed the doubled `ket * conj(bra)` tensor first, scaling as χ^(2·degree) rather than the χ^(degree+1) a good order reaches. A belief-propagation message update on a degree-3 vertex paid that cost on every sweep. For a computed order, `contract_network` now expands each operand into its contraction leaves before ordering, so `lazy(ket) * lazy(conj(bra))` reaches the optimizer as two leaves that share the physical index. The optimizer orders the full leaf set and interleaves the messages between the layers. The expansion runs on the operand after promotion to the network's common type (building on https://github.com/ITensor/ITensorNetworksNext.jl/pull/158), so the operator and state semantics are preserved, and it recurses through nested products. An explicit `order` is honored over the operands as given, so it is not flattened. The Greedy optimizer and the rest of the order machinery are unchanged. A regression test asserts the computed order's peak intermediate stays at the χ^(degree+1) path rather than forming the χ^(2·degree) doubled tensor. --- Project.toml | 2 +- src/contract_network.jl | 37 ++++++++++++++++++----- test/Project.toml | 2 ++ test/test_contract_network.jl | 55 +++++++++++++++++++++++++++++++++-- 4 files changed, 84 insertions(+), 12 deletions(-) diff --git a/Project.toml b/Project.toml index dce5a90..6680757 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "ITensorNetworksNext" uuid = "302f2e75-49f0-4526-aef7-d8ba550cb06c" -version = "0.9.12" +version = "0.9.13" authors = ["ITensor developers and contributors"] [workspace] diff --git a/src/contract_network.jl b/src/contract_network.jl index a48bdda..1e8f528 100644 --- a/src/contract_network.jl +++ b/src/contract_network.jl @@ -1,7 +1,7 @@ using Base.Broadcast: materialize using Base: @kwdef -using ITensorBase: EvaluationOrderAlgorithm, Greedy, Mul, lazy, optimize_evaluation_order, - substitute, symnameddims +using ITensorBase: EvaluationOrderAlgorithm, Greedy, Mul, ismul, lazy, + optimize_evaluation_order, substitute, symnameddims, to_mul_arguments # `contract_network` @kwdef struct Exact{Order, OrderAlg} @@ -29,19 +29,40 @@ function get_order(alg::Exact, tn) subs = Dict(symnameddims(i) => symnameddims(i, Tuple(axes(t))) for (i, t) in pairs(tn)) return substitute(order, subs) end +# The contraction leaves of an operand: a lazy product (the `NormNetwork` doubled vertex +# `lazy(ket) * lazy(conj(bra))`) contributes each factor, recursively; anything else is one leaf. +leaf_tensors(t) = ismul(t) ? mapreduce(leaf_tensors, vcat, to_mul_arguments(t)) : [t] + # 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. +# +# For a computed order, expand each promoted operand into its contraction leaves so the order +# optimizer sees each factor of a lazy product — and the physical index shared between the ket and +# bra of a doubled vertex — instead of one opaque node with only the outer bond legs. Without this +# the optimizer cannot interleave the other operands between the two layers and is forced to form +# the doubled `ket * conj(bra)` tensor first (χ^(2·degree)). Flattening the *promoted* operand keeps +# the operator/state semantics the promotion just established. An explicit operand-level `order` is +# honored over the operands as given, so it is not flattened. function contract_network(alg::Exact, tn) - order = get_order(alg, tn) + if !isnothing(alg.order) + # Explicit order: honor it over the operands as given, so it is not flattened. + order = get_order(alg, tn) + T = mapreduce(typeof, promote_type, tn) + syms_to_ts = Dict( + symnameddims(i, Tuple(axes(t))) => lazy(convert(T, t)) for (i, t) in pairs(tn) + ) + return materialize(substitute(order, syms_to_ts)) + end + # Computed order: expand each promoted operand into its contraction leaves. T = mapreduce(typeof, promote_type, tn) - syms_to_ts = Dict( - symnameddims(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) + leaves = collect(Iterators.flatten(leaf_tensors(lazy(convert(T, t))) for t in tn)) + order = get_order(alg, leaves) + syms_to_ts = + Dict(symnameddims(i, Tuple(axes(t))) => lazy(t) for (i, t) in pairs(leaves)) + return materialize(substitute(order, syms_to_ts)) end # `contraction_order` diff --git a/test/Project.toml b/test/Project.toml index a49dd6d..b83ab06 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -20,6 +20,7 @@ StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb" TensorAlgebra = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a" TensorKitSectors = "13a9c161-d5da-41f0-bcbd-e1a08ae0647f" +TermInterface = "8ea1fca8-c5ef-4a55-8b96-4e9afe9c9a3c" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [sources.ITensorNetworksNext] @@ -46,4 +47,5 @@ StableRNGs = "1" Suppressor = "0.2.8" TensorAlgebra = "0.16, 0.17" TensorKitSectors = "0.3" +TermInterface = "2" Test = "1.10" diff --git a/test/test_contract_network.jl b/test/test_contract_network.jl index e567a62..1804df4 100644 --- a/test/test_contract_network.jl +++ b/test/test_contract_network.jl @@ -1,11 +1,12 @@ using Graphs: edges, vertices -using ITensorBase: - Greedy, Index, NamedTensorOperator, inputnames, operator, outputnames, state +using ITensorBase: Greedy, ITensor, Index, NamedTensorOperator, conj, dimnames, inputnames, + lazy, operator, outputnames, state using ITensorNetworksNext: Exact, ITensorNetwork, LeftAssociative, contract_network, - linkinds, siteinds, tensornetwork + get_order, leaf_tensors, linkinds, siteinds, tensornetwork using NamedGraphs.GraphsExtensions: arranged_edges, incident_edges using NamedGraphs.NamedGraphGenerators: named_grid using OMEinsumContractionOrders: ExhaustiveSearch, GreedyMethod, TreeSA +using TermInterface: arguments, iscall using Test: @test, @testset @testset "contract_network" begin @@ -91,4 +92,52 @@ using Test: @test, @testset @test outputnames(rb) == outputnames((op * u) * w) == outputnames(op * (u * w)) @test inputnames(rb) == inputnames((op * u) * w) == inputnames(op * (u * w)) end + + @testset "Flatten lazy product operands into contraction leaves" begin + # A NormNetwork's doubled vertex is a lazy product `lazy(ket) * lazy(conj(bra))`. The + # contraction order must see each factor — and the physical index they share — as its own + # leaf, so the optimizer can interleave the other operands between the two layers instead of + # forming the doubled `ket * conj(bra)` tensor (χ^(2·degree)). + d, χ = 2, 8 + p = Index(d) + b1, b2, b3 = Index(χ), Index(χ), Index(χ) + c1, c2, c3 = Index(χ), Index(χ), Index(χ) + A = ITensor(randn(d, χ, χ, χ), (p, b1, b2, b3)) + B = ITensor(randn(d, χ, χ, χ), (p, c1, c2, c3)) + factor = lazy(A) * lazy(conj(B)) + + # `leaf_tensors` splits a lazy product into its factors, recursively; anything else is one leaf. + @test length(leaf_tensors(factor)) == 2 + @test length(leaf_tensors(lazy(A))) == 1 + @test length(leaf_tensors(lazy(factor) * lazy(A))) == 3 + + # Degree-3 doubled vertex with two incoming bond messages, contracted to the outgoing message. + msg1 = ITensor(randn(χ, χ), (b1, c1)) + msg2 = ITensor(randn(χ, χ), (b2, c2)) + net = [msg1, msg2, factor] + + # The network flattens to degree + 1 = 4 leaves: the ket and bra are separate contractible nodes. + T = mapreduce(typeof, promote_type, net) + leaves = collect(Iterators.flatten(leaf_tensors(lazy(convert(T, t))) for t in net)) + @test length(leaves) == 4 + + # Walk the computed order: a contraction node's open dims are the symmetric difference of its + # children's, its width their size product. The peak stays at the χ^(degree+1) path — no + # intermediate spans all 2·degree bond legs (the χ^(2·degree) doubled tensor). + order = get_order(Exact(), leaves) + idxsize = + Dict(n => s for t in (A, B, msg1, msg2) for (n, s) in zip(dimnames(t), size(t))) + width(dims) = isempty(dims) ? 1 : prod(idxsize[n] for n in dims) + function peakwidth(node) + iscall(node) || return (Set(dimnames(node)), width(dimnames(node))) + (da, pa), (db, pb) = peakwidth.(arguments(node)) + open = symdiff(da, db) + return open, max(pa, pb, width(open)) + end + @test last(peakwidth(order)) ≤ χ^(3 + 1) + @test last(peakwidth(order)) < χ^(2 * 3) + + # The flattened contraction matches forming the doubled vertex explicitly. + @test contract_network(net) ≈ contract_network([msg1, msg2, A, conj(B)]) + end end