Flatten lazy product operands into leaves for contraction ordering

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 <support@itensor.org> and contributors"]
 
 [workspace]

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`

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"

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