Add OMEinsumContractionOrders extension for contraction-order optimization

Project.toml

@@ -1,6 +1,6 @@
 name = "ITensorBase"
 uuid = "4795dd04-0d67-49bb-8f44-b89c448a1dc7"
-version = "0.10.5"
+version = "0.10.6"
 authors = ["ITensor developers <support@itensor.org> and contributors"]
 
 [workspace]
@@ -28,12 +28,14 @@ WrappedUnions = "325db55a-9c6c-5b90-b1a2-ec87e7a38c44"
 [weakdeps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
+OMEinsumContractionOrders = "6f22d1fd-8eed-4bb7-9776-e7d684900715"
 TensorKit = "07d1fe3e-3e46-537d-9eac-e9e13d0d4cec"
 TensorOperations = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
 
 [extensions]
 ITensorBaseAdaptExt = "Adapt"
 ITensorBaseMooncakeExt = "Mooncake"
+ITensorBaseOMEinsumContractionOrdersExt = "OMEinsumContractionOrders"
 ITensorBaseTensorKitExt = "TensorKit"
 ITensorBaseTensorOperationsExt = "TensorOperations"
 
@@ -48,6 +50,7 @@ ConstructionBase = "1.6"
 LinearAlgebra = "1.10"
 MatrixAlgebraKit = "0.2, 0.3, 0.4, 0.5, 0.6"
 Mooncake = "0.4.202, 0.5"
+OMEinsumContractionOrders = "1"
 OrderedCollections = "1.6"
 Random = "1.10"
 SimpleTraits = "0.9.4"

ext/ITensorBaseOMEinsumContractionOrdersExt/ITensorBaseOMEinsumContractionOrdersExt.jl

@@ -0,0 +1,34 @@
+module ITensorBaseOMEinsumContractionOrdersExt
+
+using ITensorBase.TermInterface: arguments
+using ITensorBase: ITensorBase, inds, ismul, optimize_contraction_order
+using OMEinsumContractionOrders:
+    OMEinsumContractionOrders, CodeOptimizer, EinCode, NestedEinsum, optimize_code
+
+# Rebuild a nested product expression from an optimized `NestedEinsum`, mapping each
+# leaf's `tensorindex` back to the corresponding argument of the flat product.
+function nested_einsum_to_expr(f, code::NestedEinsum)
+    # A leaf holds the 1-based index of its input tensor; internal nodes hold `-1`.
+    return if code.tensorindex != -1
+        f(code.tensorindex)
+    else
+        prod(Base.Fix1(nested_einsum_to_expr, f), code.args)
+    end
+end
+
+# Find a contraction order with any OMEinsumContractionOrders optimizer (`GreedyMethod`,
+# `TreeSA`, `KaHyParBipartite`, ...) by forwarding to `optimize_code`.
+function ITensorBase.optimize_contraction_order(alg::CodeOptimizer, a)
+    @assert ismul(a)
+    ts = arguments(a)
+    ixs = collect.(inds.(ts))
+    all_inds = reduce(vcat, ixs)
+    labels = unique(all_inds)
+    size_dict = Dict(i => length(i) for i in labels)
+    # 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_expr(i -> ts[i], code)
+end
+
+end

test/Project.toml

@@ -9,6 +9,7 @@ JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MatrixAlgebraKit = "6c742aac-3347-4629-af66-fc926824e5e4"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
+OMEinsumContractionOrders = "6f22d1fd-8eed-4bb7-9776-e7d684900715"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
@@ -36,6 +37,7 @@ JLArrays = "0.2, 0.3"
 LinearAlgebra = "1.10"
 MatrixAlgebraKit = "0.2, 0.3, 0.4, 0.5, 0.6"
 Mooncake = "0.4, 0.5"
+OMEinsumContractionOrders = "1"
 Random = "1.10"
 SafeTestsets = "0.1"
 StableRNGs = "1"

test/test_lazyitensors.jl

@@ -3,6 +3,7 @@ using Base.Broadcast: materialize
 using ITensorBase: @names, Greedy, LazyNamedTensor, Mul, NamedTensor, Optimal,
     SymbolicNamedTensor, dimnames, inds, ismul, lazy, nameddims, namedoneto,
     optimize_evaluation_order, substitute, symnameddims
+using OMEinsumContractionOrders: GreedyMethod, TreeSA
 using TensorOperations: TensorOperations
 using TermInterface: arguments, arity, children, head, iscall, isexpr, maketerm, operation,
     sorted_arguments, sorted_children
@@ -119,4 +120,19 @@ using WrappedUnions: unwrap
         @test arity(ordered) == 2
         @test issetequal(dimnames(ordered), dimnames(flat))
     end
+
+    @testset "optimize_evaluation_order (OMEinsumContractionOrders $alg)" for alg in
+        (
+            GreedyMethod(),
+            TreeSA(),
+        )
+        i, j, k, l = namedoneto.((2, 3, 4, 5), (:i, :j, :k, :l))
+        s = [symnameddims(:a, (i, j)), symnameddims(:b, (j, k)), symnameddims(:c, (k, l))]
+        flat = lazy(Mul(s))
+        ordered = optimize_evaluation_order(flat; alg)
+        @test ordered isa LazyNamedTensor
+        @test ismul(ordered)
+        @test arity(ordered) == 2
+        @test issetequal(dimnames(ordered), dimnames(flat))
+    end
 end