From 4e9482493217c718e1f1adc714a5bfc55b6e0e5d Mon Sep 17 00:00:00 2001 From: Matthew Fishman Date: Thu, 2 Jul 2026 18:29:12 -0400 Subject: [PATCH] Add OMEinsumContractionOrders extension for contraction-order optimization Adds an OMEinsumContractionOrders extension so any of its optimizers (`GreedyMethod`, `TreeSA`, `KaHyParBipartite`, `SABipartite`, `ExactTreewidth`, ...) can be passed directly as the algorithm to `optimize_evaluation_order`/`optimize_contraction_order`. The extension dispatches on the abstract `CodeOptimizer`, so it forwards whichever optimizers the resolved OMEinsumContractionOrders version provides (including `ExhaustiveSearch` once released) with no per-optimizer wrapping. It builds an `EinCode` from the flattened product, calls `optimize_code`, and translates the returned `NestedEinsum` back into a nested contraction expression. --- Project.toml | 5 ++- ...ITensorBaseOMEinsumContractionOrdersExt.jl | 34 +++++++++++++++++++ test/Project.toml | 2 ++ test/test_lazyitensors.jl | 16 +++++++++ 4 files changed, 56 insertions(+), 1 deletion(-) create mode 100644 ext/ITensorBaseOMEinsumContractionOrdersExt/ITensorBaseOMEinsumContractionOrdersExt.jl diff --git a/Project.toml b/Project.toml index 383cc22..f4aea5f 100644 --- a/Project.toml +++ b/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 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" diff --git a/ext/ITensorBaseOMEinsumContractionOrdersExt/ITensorBaseOMEinsumContractionOrdersExt.jl b/ext/ITensorBaseOMEinsumContractionOrdersExt/ITensorBaseOMEinsumContractionOrdersExt.jl new file mode 100644 index 0000000..048bc1c --- /dev/null +++ b/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 diff --git a/test/Project.toml b/test/Project.toml index 882bd69..5a2e453 100644 --- a/test/Project.toml +++ b/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" diff --git a/test/test_lazyitensors.jl b/test/test_lazyitensors.jl index 1bca384..e4a5e9b 100644 --- a/test/test_lazyitensors.jl +++ b/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