From e4b102db34b78d6d73e83bd2df0dfda942c7d7f7 Mon Sep 17 00:00:00 2001
From: ArrogantGao <gaoxuanzhao@gmail.com>
Date: Thu, 1 Aug 2024 20:47:26 +0800
Subject: [PATCH 01/15] add OMEinsumContractionOrders as backend of optimaltree

---
 .gitignore                                    |  4 +-
 Project.toml                                  | 21 ++---
 ...rOperationsOMEinsumContractionOrdersExt.jl | 79 +++++++++++++++++++
 src/TensorOperations.jl                       |  3 +
 src/indexnotation/optimaltree.jl              | 13 ++-
 src/indexnotation/tensormacros.jl             |  9 ++-
 6 files changed, 110 insertions(+), 19 deletions(-)
 create mode 100644 ext/TensorOperationsOMEinsumContractionOrdersExt.jl

diff --git a/.gitignore b/.gitignore
index 0ee3d176..5caf9ac5 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,4 +1,6 @@
 *.jl.cov
 *.jl.*.cov
 *.jl.mem
-Manifest.toml
\ No newline at end of file
+Manifest.toml
+.vscode
+.DS_Store
\ No newline at end of file
diff --git a/Project.toml b/Project.toml
index 8c8d281c..9a1d84c9 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,10 +1,6 @@
 name = "TensorOperations"
 uuid = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
-authors = [
-    "Lukas Devos <lukas.devos@ugent.be>",
-    "Maarten Van Damme <maartenvd1994@gmail.com>",
-    "Jutho Haegeman <jutho.haegeman@ugent.be>",
-]
+authors = ["Lukas Devos <lukas.devos@ugent.be>", "Maarten Van Damme <maartenvd1994@gmail.com>", "Jutho Haegeman <jutho.haegeman@ugent.be>"]
 version = "5.0.0"
 
 [deps]
@@ -25,10 +21,12 @@ Bumper = "8ce10254-0962-460f-a3d8-1f77fea1446e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1"
+OMEinsumContractionOrders = "6f22d1fd-8eed-4bb7-9776-e7d684900715"
 
 [extensions]
 TensorOperationsBumperExt = "Bumper"
 TensorOperationsChainRulesCoreExt = "ChainRulesCore"
+TensorOperationsOMEinsumContractionOrdersExt = "OMEinsumContractionOrders"
 TensorOperationscuTENSORExt = ["cuTENSOR", "CUDA"]
 
 [compat]
@@ -41,6 +39,7 @@ DynamicPolynomials = "0.5"
 LRUCache = "1"
 LinearAlgebra = "1.6"
 Logging = "1.6"
+OMEinsumContractionOrders = "0.8.3"
 PackageExtensionCompat = "1"
 PtrArrays = "1.2"
 Random = "1"
@@ -64,14 +63,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1"
 
 [targets]
-test = [
-    "Test",
-    "Random",
-    "DynamicPolynomials",
-    "ChainRulesTestUtils",
-    "CUDA",
-    "cuTENSOR",
-    "Aqua",
-    "Logging",
-    "Bumper",
-]
+test = ["Test", "Random", "DynamicPolynomials", "ChainRulesTestUtils", "CUDA", "cuTENSOR", "Aqua", "Logging", "Bumper"]
diff --git a/ext/TensorOperationsOMEinsumContractionOrdersExt.jl b/ext/TensorOperationsOMEinsumContractionOrdersExt.jl
new file mode 100644
index 00000000..be63f2a1
--- /dev/null
+++ b/ext/TensorOperationsOMEinsumContractionOrdersExt.jl
@@ -0,0 +1,79 @@
+module TensorOperationsOMEinsumContractionOrdersExt
+
+using TensorOperations
+using TensorOperations: TensorOperations as TO
+using TensorOperations: TreeOptimizer
+using OMEinsumContractionOrders
+using OMEinsumContractionOrders: EinCode, NestedEinsum, SlicedEinsum, isleaf
+
+function TO.optimaltree(network, optdata::Dict{TDK, TDV}, ::TreeOptimizer{:GreedyMethod}, verbose::Bool) where{TDK, TDV}
+    ome_optimizer = GreedyMethod()
+    return optimize(network, optdata, ome_optimizer, verbose)
+end
+
+function TO.optimaltree(network, optdata::Dict{TDK, TDV}, ::TreeOptimizer{:KaHyParBipartite}, verbose::Bool) where{TDK, TDV}
+    # sc_target and max_group_size are simply set as 10 for now
+    ome_optimizer = KaHyParBipartite(; sc_target=10, max_group_size=10)
+    return optimize(network, optdata, ome_optimizer, verbose)
+end
+
+function TO.optimaltree(network, optdata::Dict{TDK, TDV}, ::TreeOptimizer{:TreeSA}, verbose::Bool) where{TDK, TDV}
+    ome_optimizer = TreeSA()
+    return optimize(network, optdata, ome_optimizer, verbose)
+end
+
+function TO.optimaltree(network, optdata::Dict{TDK, TDV}, ::TreeOptimizer{:SABipartite}, verbose::Bool) where{TDK, TDV}
+    ome_optimizer = SABipartite()
+    return optimize(network, optdata, ome_optimizer, verbose)
+end
+
+function optimize(network, optdata::Dict{TDK, TDV}, ome_optimizer::CodeOptimizer, verbose::Bool) where{TDK, TDV}
+
+    try
+        TDV <: Number
+    catch
+        error("Dimensions of the tensors must be numbers")
+    end
+
+    # transform the network as EinCode
+    code = network2eincode(network)
+    # optimize the contraction order using OMEinsumContractionOrders, which gives a NestedEinsum
+    optcode = optimize_code(code, optdata, ome_optimizer)
+
+    if verbose
+        cc = OMEinsumContractionOrders.contraction_complexity(optcode, optdata)
+        println("Optimized contraction order: ", optcode)
+        println("Contraction complexity: ", cc)
+    end
+
+    # transform the optimized contraction order back to the network
+    optimaltree = eincode2contractiontree(optcode)
+
+    return (optimaltree,)
+end
+
+function network2eincode(network)
+    indices = unique(vcat(network...))
+    open_edges = empty(network[1])
+    # if a indices appear only once, it is an open index
+    for i in indices
+        if sum([i in t for t in network]) == 1
+            push!(open_edges, i)
+        end
+    end
+    return EinCode(network, open_edges)
+end
+
+function eincode2contractiontree(eincode::NestedEinsum)
+    if isleaf(eincode)
+        return eincode.tensorindex
+    else
+        return [eincode2contractiontree(arg) for arg in eincode.args]
+    end
+end
+
+function eincode2contractiontree(eincode::SlicedEinsum)
+    return eincode2contractiontree(eincode.eins)
+end
+
+end
\ No newline at end of file
diff --git a/src/TensorOperations.jl b/src/TensorOperations.jl
index d3e6ff7d..7f3aa8d4 100644
--- a/src/TensorOperations.jl
+++ b/src/TensorOperations.jl
@@ -24,6 +24,9 @@ export tensorcopy!, tensoradd!, tensortrace!, tensorcontract!, tensorproduct!, t
 export tensorcopy, tensoradd, tensortrace, tensorcontract, tensorproduct, scalartype
 export tensoralloc, tensorfree!
 
+# export optimizer API
+export DefaultOptimizer, OMEinsumOptimizer
+
 export IndexTuple, Index2Tuple, linearize
 
 # export debug functionality
diff --git a/src/indexnotation/optimaltree.jl b/src/indexnotation/optimaltree.jl
index 5e1fd101..62c4d49c 100644
--- a/src/indexnotation/optimaltree.jl
+++ b/src/indexnotation/optimaltree.jl
@@ -1,4 +1,15 @@
-function optimaltree(network, optdata::Dict; verbose::Bool=false)
+struct TreeOptimizer{T} end # T is a Symbol for the algorithm
+
+function optimaltree(network, optdata::Dict; optimizer::TreeOptimizer = TreeOptimizer{:NCon}(), verbose::Bool=false)
+    @debug "Using optimizer $(typeof(optimizer))"
+    return optimaltree(network, optdata, optimizer, verbose)
+end
+
+function optimaltree(network, optdata::Dict, ::TreeOptimizer{T}, verbose::Bool) where{T}
+    throw(ArgumentError("Unknown optimizer: $T"))
+end
+
+function optimaltree(network, optdata::Dict, ::TreeOptimizer{:NCon}, verbose::Bool)
     numtensors = length(network)
     allindices = unique(vcat(network...))
     numindices = length(allindices)
diff --git a/src/indexnotation/tensormacros.jl b/src/indexnotation/tensormacros.jl
index 678aa983..4581e2dd 100644
--- a/src/indexnotation/tensormacros.jl
+++ b/src/indexnotation/tensormacros.jl
@@ -71,6 +71,7 @@ function tensorparser(tensorexpr, kwargs...)
         end
     end
     # now handle the remaining keyword arguments
+    optimizer = TreeOptimizer{:NCon}() # the default optimizer
     for (name, val) in kwargs
         if name == :order
             isexpr(val, :tuple) ||
@@ -85,6 +86,12 @@ function tensorparser(tensorexpr, kwargs...)
             val in (:warn, :cache) ||
                 throw(ArgumentError("Invalid use of `costcheck`, should be `costcheck=warn` or `costcheck=cache`"))
             parser.contractioncostcheck = val
+        elseif name == :opt_algorithm
+            if val isa Symbol
+                optimizer = TreeOptimizer{val}()
+            else
+                throw(ArgumentError("Invalid use of `opt_algorithm`, should be `opt_algorithm=NCon` or `opt_algorithm=NameOfAlgorithm`"))
+            end
         elseif name == :opt
             if val isa Bool && val
                 optdict = optdata(tensorexpr)
@@ -93,7 +100,7 @@ function tensorparser(tensorexpr, kwargs...)
             else
                 throw(ArgumentError("Invalid use of `opt`, should be `opt=true` or `opt=OptExpr`"))
             end
-            parser.contractiontreebuilder = network -> optimaltree(network, optdict)[1]
+            parser.contractiontreebuilder = network -> optimaltree(network, optdict, optimizer = optimizer)[1]
         elseif !(name == :backend || name == :allocator) # these two have been handled
             throw(ArgumentError("Unknown keyword argument `name`."))
         end

From d6bbe6fe1e54419f4ae4aae51bd9b1880e55994b Mon Sep 17 00:00:00 2001
From: ArrogantGao <gaoxuanzhao@gmail.com>
Date: Thu, 1 Aug 2024 22:05:43 +0800
Subject: [PATCH 02/15] fixed tests

---
 Project.toml                                  |   3 +-
 ...rOperationsOMEinsumContractionOrdersExt.jl |   1 +
 src/TensorOperations.jl                       |   3 -
 test/macro_kwargs.jl                          |  16 +++
 test/omeinsumcontractionordres.jl             | 100 ++++++++++++++++++
 test/runtests.jl                              |  72 +++++++------
 6 files changed, 158 insertions(+), 37 deletions(-)
 create mode 100644 test/omeinsumcontractionordres.jl

diff --git a/Project.toml b/Project.toml
index 9a1d84c9..6f79ed25 100644
--- a/Project.toml
+++ b/Project.toml
@@ -61,6 +61,7 @@ Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1"
+OMEinsumContractionOrders = "6f22d1fd-8eed-4bb7-9776-e7d684900715"
 
 [targets]
-test = ["Test", "Random", "DynamicPolynomials", "ChainRulesTestUtils", "CUDA", "cuTENSOR", "Aqua", "Logging", "Bumper"]
+test = ["Test", "Random", "DynamicPolynomials", "ChainRulesTestUtils", "CUDA", "cuTENSOR", "Aqua", "Logging", "Bumper", "OMEinsumContractionOrders"]
diff --git a/ext/TensorOperationsOMEinsumContractionOrdersExt.jl b/ext/TensorOperationsOMEinsumContractionOrdersExt.jl
index be63f2a1..24a80a30 100644
--- a/ext/TensorOperationsOMEinsumContractionOrdersExt.jl
+++ b/ext/TensorOperationsOMEinsumContractionOrdersExt.jl
@@ -72,6 +72,7 @@ function eincode2contractiontree(eincode::NestedEinsum)
     end
 end
 
+# TreeSA returns a SlicedEinsum, with nslice = 0, so directly using the eins
 function eincode2contractiontree(eincode::SlicedEinsum)
     return eincode2contractiontree(eincode.eins)
 end
diff --git a/src/TensorOperations.jl b/src/TensorOperations.jl
index 7f3aa8d4..d3e6ff7d 100644
--- a/src/TensorOperations.jl
+++ b/src/TensorOperations.jl
@@ -24,9 +24,6 @@ export tensorcopy!, tensoradd!, tensortrace!, tensorcontract!, tensorproduct!, t
 export tensorcopy, tensoradd, tensortrace, tensorcontract, tensorproduct, scalartype
 export tensoralloc, tensorfree!
 
-# export optimizer API
-export DefaultOptimizer, OMEinsumOptimizer
-
 export IndexTuple, Index2Tuple, linearize
 
 # export debug functionality
diff --git a/test/macro_kwargs.jl b/test/macro_kwargs.jl
index 254b6a8f..64826d9a 100644
--- a/test/macro_kwargs.jl
+++ b/test/macro_kwargs.jl
@@ -106,3 +106,19 @@ end
     end
     @test D1 ≈ D2 ≈ D3 ≈ D4 ≈ D5
 end
+
+@testset "opt_algorithm" begin
+    A = randn(5, 5, 5, 5)
+    B = randn(5, 5, 5)
+    C = randn(5, 5, 5)
+
+    @tensor opt = true begin
+        D1[a, b, c, d] := A[a, e, c, f] * B[g, d, e] * C[g, f, b]
+    end
+
+    @tensor opt = true opt_algorithm = NCon begin
+        D2[a, b, c, d] := A[a, e, c, f] * B[g, d, e] * C[g, f, b]
+    end
+
+    @test D1 ≈ D2
+end
\ No newline at end of file
diff --git a/test/omeinsumcontractionordres.jl b/test/omeinsumcontractionordres.jl
new file mode 100644
index 00000000..6292c292
--- /dev/null
+++ b/test/omeinsumcontractionordres.jl
@@ -0,0 +1,100 @@
+@testset "@tensor dependency check" begin
+    @test_throws ArgumentError begin
+        A = rand(2, 2)
+        B = rand(2, 2)
+        C = rand(2, 2)
+        ex = :(@tensor opt=(i=>2, j=>2, k=>2) opt_algorithm=GreedyMethod S[] := A[i, j] * B[j, k] * C[i, k])
+        macroexpand(Main, ex)
+    end
+end
+
+using OMEinsumContractionOrders
+
+@testset "OMEinsumContractionOrders optimization algorithms" begin
+    A = randn(5, 5, 5, 5)
+    B = randn(5, 5, 5)
+    C = randn(5, 5, 5)
+
+    @tensor begin
+        D1[a, b, c, d] := A[a, e, c, f] * B[g, d, e] * C[g, f, b]
+    end
+
+    @tensor opt = (a => 5, b => 5, c => 5, d => 5, e => 5, f => 5, g => 5) opt_algorithm = GreedyMethod begin
+        D2[a, b, c, d] := A[a, e, c, f] * B[g, d, e] * C[g, f, b]
+    end
+
+    @tensor opt = (a => 5, b => 5, c => 5, d => 5, e => 5, f => 5, g => 5) opt_algorithm = TreeSA begin
+        D3[a, b, c, d] := A[a, e, c, f] * B[g, d, e] * C[g, f, b]
+    end
+
+    @tensor opt = (a => 5, b => 5, c => 5, d => 5, e => 5, f => 5, g => 5) opt_algorithm = KaHyParBipartite begin
+        D4[a, b, c, d] := A[a, e, c, f] * B[g, d, e] * C[g, f, b]
+    end
+
+    @tensor opt = (a => 5, b => 5, c => 5, d => 5, e => 5, f => 5, g => 5) opt_algorithm = SABipartite begin
+        D5[a, b, c, d] := A[a, e, c, f] * B[g, d, e] * C[g, f, b]
+    end
+
+    @tensor opt = (1 => 5, 2 => 5, 3 => 5, 4 => 5, 5 => 5, 6 => 5, 7 => 5) opt_algorithm = GreedyMethod begin
+        D6[1, 2, 3, 4] := A[1, 5, 3, 6] * B[7, 4, 5] * C[7, 6, 2]
+    end
+
+    @test D1 ≈ D2 ≈ D3 ≈ D4 ≈ D5 ≈ D6
+
+
+    A = rand(2, 2)
+    B = rand(2, 2, 2)
+    C = rand(2, 2)
+    D = rand(2, 2)
+    E = rand(2, 2, 2)
+    F = rand(2, 2)
+    
+    @tensor opt = true begin
+        s1[] := A[i, k] * B[i, j, l] * C[j, m] * D[k, n] * E[n, l, o] * F[o, m]
+    end
+
+    @tensor opt = (i => 2, j => 2, k => 2, l => 2, m => 2, n => 2, o => 2) opt_algorithm = GreedyMethod begin
+        s2[] := A[i, k] * B[i, j, l] * C[j, m] * D[k, n] * E[n, l, o] * F[o, m]
+    end
+
+    @tensor opt = (i => 2, j => 2, k => 2, l => 2, m => 2, n => 2, o => 2) opt_algorithm = TreeSA begin
+        s3[] := A[i, k] * B[i, j, l] * C[j, m] * D[k, n] * E[n, l, o] * F[o, m]
+    end
+
+    @tensor opt = (i => 2, j => 2, k => 2, l => 2, m => 2, n => 2, o => 2) opt_algorithm = KaHyParBipartite begin
+        s4[] := A[i, k] * B[i, j, l] * C[j, m] * D[k, n] * E[n, l, o] * F[o, m]
+    end
+
+    @tensor opt = (i => 2, j => 2, k => 2, l => 2, m => 2, n => 2, o => 2) opt_algorithm = SABipartite begin
+        s5[] := A[i, k] * B[i, j, l] * C[j, m] * D[k, n] * E[n, l, o] * F[o, m]
+    end
+
+    @test s1 ≈ s2 ≈ s3 ≈ s4 ≈ s5
+
+    A = randn(5, 5, 5)
+    B = randn(5, 5, 5)
+    C = randn(5, 5, 5)
+    α = randn()
+
+    @tensor opt = true begin
+        D1[m] := A[i, j, k] * B[j, k, l] * C[i, l, m] + α * A[i, j, k] * B[j, k, l] * C[i, l, m]
+    end
+
+    @tensor opt = (i => 5, j => 5, k => 5, l => 5, m => 5) opt_algorithm = GreedyMethod begin
+        D2[m] := A[i, j, k] * B[j, k, l] * C[i, l, m] + α * A[i, j, k] * B[j, k, l] * C[i, l, m]
+    end
+
+    @tensor opt = (i => 5, j => 5, k => 5, l => 5, m => 5) opt_algorithm = TreeSA begin
+        D3[m] := A[i, j, k] * B[j, k, l] * C[i, l, m] + α * A[i, j, k] * B[j, k, l] * C[i, l, m]
+    end
+
+    @tensor opt = (i => 5, j => 5, k => 5, l => 5, m => 5) opt_algorithm = KaHyParBipartite begin
+        D4[m] := A[i, j, k] * B[j, k, l] * C[i, l, m] + α * A[i, j, k] * B[j, k, l] * C[i, l, m]
+    end
+
+    @tensor opt = (i => 5, j => 5, k => 5, l => 5, m => 5) opt_algorithm = SABipartite begin
+        D5[m] := A[i, j, k] * B[j, k, l] * C[i, l, m] + α * A[i, j, k] * B[j, k, l] * C[i, l, m]
+    end
+
+    @test D1 ≈ D2 ≈ D3 ≈ D4 ≈ D5
+end
\ No newline at end of file
diff --git a/test/runtests.jl b/test/runtests.jl
index 5b068140..3b601729 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -8,42 +8,48 @@ using TensorOperations: IndexError
 using TensorOperations: BaseCopy, BaseView, StridedNative, StridedBLAS
 using TensorOperations: DefaultAllocator, ManualAllocator
 
-@testset "tensoropt" verbose = true begin
-    include("tensoropt.jl")
-end
-@testset "auxiliary" verbose = true begin
-    include("auxiliary.jl")
-end
-@testset "macro keywords" verbose = true begin
-    include("macro_kwargs.jl")
-end
-@testset "method syntax" verbose = true begin
-    include("methods.jl")
-end
-@testset "macro with index notation" verbose = true begin
-    include("tensor.jl")
-end
-@testset "ad" verbose = false begin
-    include("ad.jl")
-end
+# @testset "tensoropt" verbose = true begin
+#     include("tensoropt.jl")
+# end
+# @testset "auxiliary" verbose = true begin
+#     include("auxiliary.jl")
+# end
+# @testset "macro keywords" verbose = true begin
+#     include("macro_kwargs.jl")
+# end
+# @testset "method syntax" verbose = true begin
+#     include("methods.jl")
+# end
+# @testset "macro with index notation" verbose = true begin
+#     include("tensor.jl")
+# end
+# @testset "ad" verbose = false begin
+#     include("ad.jl")
+# end
 
-# note: cuTENSOR should not be loaded before this point
-# as there is a test which requires it to be loaded after
-@testset "cuTENSOR extension" verbose = true begin
-    include("cutensor.jl")
-end
+# # note: cuTENSOR should not be loaded before this point
+# # as there is a test which requires it to be loaded after
+# @testset "cuTENSOR extension" verbose = true begin
+#     include("cutensor.jl")
+# end
+
+# # note: Bumper should not be loaded before this point
+# # as there is a test which requires it to be loaded after
+# @testset "Bumper extension" verbose = true begin
+#     include("butensor.jl")
+# end
 
-# note: Bumper should not be loaded before this point
+# note: OMEinsumContractionOrders should not be loaded before this point
 # as there is a test which requires it to be loaded after
-@testset "Bumper extension" verbose = true begin
-    include("butensor.jl")
+@testset "OMEinsumOptimizer extension" begin
+    include("omeinsumcontractionordres.jl")
 end
 
-@testset "Polynomials" begin
-    include("polynomials.jl")
-end
+# @testset "Polynomials" begin
+#     include("polynomials.jl")
+# end
 
-@testset "Aqua" verbose = true begin
-    using Aqua
-    Aqua.test_all(TensorOperations)
-end
+# @testset "Aqua" verbose = true begin
+#     using Aqua
+#     Aqua.test_all(TensorOperations)
+# end

From 6d611030d274d88859b96dd53c881d6fef603623 Mon Sep 17 00:00:00 2001
From: ArrogantGao <gaoxuanzhao@gmail.com>
Date: Thu, 1 Aug 2024 22:07:10 +0800
Subject: [PATCH 03/15] fixed size dict must be Dict{T, Number}

---
 ext/TensorOperationsOMEinsumContractionOrdersExt.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ext/TensorOperationsOMEinsumContractionOrdersExt.jl b/ext/TensorOperationsOMEinsumContractionOrdersExt.jl
index 24a80a30..0f272c35 100644
--- a/ext/TensorOperationsOMEinsumContractionOrdersExt.jl
+++ b/ext/TensorOperationsOMEinsumContractionOrdersExt.jl
@@ -30,7 +30,7 @@ end
 function optimize(network, optdata::Dict{TDK, TDV}, ome_optimizer::CodeOptimizer, verbose::Bool) where{TDK, TDV}
 
     try
-        TDV <: Number
+        @assert TDV <: Number
     catch
         error("Dimensions of the tensors must be numbers")
     end

From 329952d13a3ce2484ce8a9b75a13fe22baea9650 Mon Sep 17 00:00:00 2001
From: ArrogantGao <gaoxuanzhao@gmail.com>
Date: Thu, 1 Aug 2024 22:07:21 +0800
Subject: [PATCH 04/15] fix optimize

---
 ext/TensorOperationsOMEinsumContractionOrdersExt.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ext/TensorOperationsOMEinsumContractionOrdersExt.jl b/ext/TensorOperationsOMEinsumContractionOrdersExt.jl
index 0f272c35..ece83498 100644
--- a/ext/TensorOperationsOMEinsumContractionOrdersExt.jl
+++ b/ext/TensorOperationsOMEinsumContractionOrdersExt.jl
@@ -32,7 +32,7 @@ function optimize(network, optdata::Dict{TDK, TDV}, ome_optimizer::CodeOptimizer
     try
         @assert TDV <: Number
     catch
-        error("Dimensions of the tensors must be numbers")
+        throw(ArgumentError("The values of the optdata dictionary must be of type Number"))
     end
 
     # transform the network as EinCode

From c14fa7c8fc88183ac20da9d5291fc44124023b8c Mon Sep 17 00:00:00 2001
From: ArrogantGao <gaoxuanzhao@gmail.com>
Date: Fri, 2 Aug 2024 20:54:17 +0800
Subject: [PATCH 05/15] update optimize, reture cost of the contraction

---
 ext/TensorOperationsOMEinsumContractionOrdersExt.jl | 12 +++++-------
 src/indexnotation/optimaltree.jl                    |  2 +-
 2 files changed, 6 insertions(+), 8 deletions(-)

diff --git a/ext/TensorOperationsOMEinsumContractionOrdersExt.jl b/ext/TensorOperationsOMEinsumContractionOrdersExt.jl
index ece83498..f80465f9 100644
--- a/ext/TensorOperationsOMEinsumContractionOrdersExt.jl
+++ b/ext/TensorOperationsOMEinsumContractionOrdersExt.jl
@@ -40,16 +40,14 @@ function optimize(network, optdata::Dict{TDK, TDV}, ome_optimizer::CodeOptimizer
     # optimize the contraction order using OMEinsumContractionOrders, which gives a NestedEinsum
     optcode = optimize_code(code, optdata, ome_optimizer)
 
-    if verbose
-        cc = OMEinsumContractionOrders.contraction_complexity(optcode, optdata)
-        println("Optimized contraction order: ", optcode)
-        println("Contraction complexity: ", cc)
-    end
-
     # transform the optimized contraction order back to the network
     optimaltree = eincode2contractiontree(optcode)
 
-    return (optimaltree,)
+    # calculate the size of maximum tensor during the contraction
+    cc = OMEinsumContractionOrders.contraction_complexity(optcode, optdata)
+    space_complexity = 2.0^cc.sc
+
+    return optimaltree, space_complexity
 end
 
 function network2eincode(network)
diff --git a/src/indexnotation/optimaltree.jl b/src/indexnotation/optimaltree.jl
index 62c4d49c..ac4048c6 100644
--- a/src/indexnotation/optimaltree.jl
+++ b/src/indexnotation/optimaltree.jl
@@ -6,7 +6,7 @@ function optimaltree(network, optdata::Dict; optimizer::TreeOptimizer = TreeOpti
 end
 
 function optimaltree(network, optdata::Dict, ::TreeOptimizer{T}, verbose::Bool) where{T}
-    throw(ArgumentError("Unknown optimizer: $T"))
+    throw(ArgumentError("Unknown optimizer: $T. Hint: may need to load extensions, e.g. `using OMEinsumContractionOrders`"))
 end
 
 function optimaltree(network, optdata::Dict, ::TreeOptimizer{:NCon}, verbose::Bool)

From 76698c195aade2b5551e45da0ef1c943e53d0702 Mon Sep 17 00:00:00 2001
From: ArrogantGao <gaoxuanzhao@gmail.com>
Date: Tue, 6 Aug 2024 10:20:49 +0800
Subject: [PATCH 06/15] add ExactTreeWidth solver

---
 Project.toml                                  |  6 +++---
 ...rOperationsOMEinsumContractionOrdersExt.jl |  5 +++++
 test/omeinsumcontractionordres.jl             | 20 +++++++++++++++----
 3 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/Project.toml b/Project.toml
index 6f79ed25..eecd27f8 100644
--- a/Project.toml
+++ b/Project.toml
@@ -20,8 +20,8 @@ cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1"
 Bumper = "8ce10254-0962-460f-a3d8-1f77fea1446e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1"
 OMEinsumContractionOrders = "6f22d1fd-8eed-4bb7-9776-e7d684900715"
+cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1"
 
 [extensions]
 TensorOperationsBumperExt = "Bumper"
@@ -39,7 +39,7 @@ DynamicPolynomials = "0.5"
 LRUCache = "1"
 LinearAlgebra = "1.6"
 Logging = "1.6"
-OMEinsumContractionOrders = "0.8.3"
+OMEinsumContractionOrders = "0.9"
 PackageExtensionCompat = "1"
 PtrArrays = "1.2"
 Random = "1"
@@ -58,10 +58,10 @@ CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 DynamicPolynomials = "7c1d4256-1411-5781-91ec-d7bc3513ac07"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
+OMEinsumContractionOrders = "6f22d1fd-8eed-4bb7-9776-e7d684900715"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1"
-OMEinsumContractionOrders = "6f22d1fd-8eed-4bb7-9776-e7d684900715"
 
 [targets]
 test = ["Test", "Random", "DynamicPolynomials", "ChainRulesTestUtils", "CUDA", "cuTENSOR", "Aqua", "Logging", "Bumper", "OMEinsumContractionOrders"]
diff --git a/ext/TensorOperationsOMEinsumContractionOrdersExt.jl b/ext/TensorOperationsOMEinsumContractionOrdersExt.jl
index f80465f9..5e270133 100644
--- a/ext/TensorOperationsOMEinsumContractionOrdersExt.jl
+++ b/ext/TensorOperationsOMEinsumContractionOrdersExt.jl
@@ -27,6 +27,11 @@ function TO.optimaltree(network, optdata::Dict{TDK, TDV}, ::TreeOptimizer{:SABip
     return optimize(network, optdata, ome_optimizer, verbose)
 end
 
+function TO.optimaltree(network, optdata::Dict{TDK, TDV}, ::TreeOptimizer{:ExactTreeWidth}, verbose::Bool) where{TDK, TDV}
+    ome_optimizer = ExactTreeWidth()
+    return optimize(network, optdata, ome_optimizer, verbose)
+end
+
 function optimize(network, optdata::Dict{TDK, TDV}, ome_optimizer::CodeOptimizer, verbose::Bool) where{TDK, TDV}
 
     try
diff --git a/test/omeinsumcontractionordres.jl b/test/omeinsumcontractionordres.jl
index 6292c292..76ea819f 100644
--- a/test/omeinsumcontractionordres.jl
+++ b/test/omeinsumcontractionordres.jl
@@ -35,11 +35,15 @@ using OMEinsumContractionOrders
         D5[a, b, c, d] := A[a, e, c, f] * B[g, d, e] * C[g, f, b]
     end
 
+    @tensor opt = (a => 5, b => 5, c => 5, d => 5, e => 5, f => 5, g => 5) opt_algorithm = ExactTreeWidth begin
+        D6[a, b, c, d] := A[a, e, c, f] * B[g, d, e] * C[g, f, b]
+    end
+
     @tensor opt = (1 => 5, 2 => 5, 3 => 5, 4 => 5, 5 => 5, 6 => 5, 7 => 5) opt_algorithm = GreedyMethod begin
-        D6[1, 2, 3, 4] := A[1, 5, 3, 6] * B[7, 4, 5] * C[7, 6, 2]
+        D7[1, 2, 3, 4] := A[1, 5, 3, 6] * B[7, 4, 5] * C[7, 6, 2]
     end
 
-    @test D1 ≈ D2 ≈ D3 ≈ D4 ≈ D5 ≈ D6
+    @test D1 ≈ D2 ≈ D3 ≈ D4 ≈ D5 ≈ D6 ≈ D7
 
 
     A = rand(2, 2)
@@ -69,7 +73,11 @@ using OMEinsumContractionOrders
         s5[] := A[i, k] * B[i, j, l] * C[j, m] * D[k, n] * E[n, l, o] * F[o, m]
     end
 
-    @test s1 ≈ s2 ≈ s3 ≈ s4 ≈ s5
+    @tensor opt = (i => 2, j => 2, k => 2, l => 2, m => 2, n => 2, o => 2) opt_algorithm = ExactTreeWidth() begin
+        s6[] := A[i, k] * B[i, j, l] * C[j, m] * D[k, n] * E[n, l, o] * F[o, m]
+    end
+
+    @test s1 ≈ s2 ≈ s3 ≈ s4 ≈ s5 ≈ s6
 
     A = randn(5, 5, 5)
     B = randn(5, 5, 5)
@@ -96,5 +104,9 @@ using OMEinsumContractionOrders
         D5[m] := A[i, j, k] * B[j, k, l] * C[i, l, m] + α * A[i, j, k] * B[j, k, l] * C[i, l, m]
     end
 
-    @test D1 ≈ D2 ≈ D3 ≈ D4 ≈ D5
+    @tensor opt = (i => 5, j => 5, k => 5, l => 5, m => 5) opt_algorithm = ExactTreeWidth begin
+        D6[m] := A[i, j, k] * B[j, k, l] * C[i, l, m] + α * A[i, j, k] * B[j, k, l] * C[i, l, m]
+    end
+
+    @test D1 ≈ D2 ≈ D3 ≈ D4 ≈ D5 ≈ D6
 end
\ No newline at end of file

From 81d64af451277e088feb69b6c72735290f212eae Mon Sep 17 00:00:00 2001
From: ArrogantGao <gaoxuanzhao@gmail.com>
Date: Wed, 7 Aug 2024 19:15:19 +0800
Subject: [PATCH 07/15] convert type of the indices as Int for ExactTreewidth
 solver

---
 ...rOperationsOMEinsumContractionOrdersExt.jl | 28 ++++----
 test/omeinsumcontractionordres.jl             |  6 +-
 test/runtests.jl                              | 70 +++++++++----------
 3 files changed, 52 insertions(+), 52 deletions(-)

diff --git a/ext/TensorOperationsOMEinsumContractionOrdersExt.jl b/ext/TensorOperationsOMEinsumContractionOrdersExt.jl
index 5e270133..437f1545 100644
--- a/ext/TensorOperationsOMEinsumContractionOrdersExt.jl
+++ b/ext/TensorOperationsOMEinsumContractionOrdersExt.jl
@@ -27,13 +27,12 @@ function TO.optimaltree(network, optdata::Dict{TDK, TDV}, ::TreeOptimizer{:SABip
     return optimize(network, optdata, ome_optimizer, verbose)
 end
 
-function TO.optimaltree(network, optdata::Dict{TDK, TDV}, ::TreeOptimizer{:ExactTreeWidth}, verbose::Bool) where{TDK, TDV}
-    ome_optimizer = ExactTreeWidth()
+function TO.optimaltree(network, optdata::Dict{TDK, TDV}, ::TreeOptimizer{:ExactTreewidth}, verbose::Bool) where{TDK, TDV}
+    ome_optimizer = ExactTreewidth()
     return optimize(network, optdata, ome_optimizer, verbose)
 end
 
 function optimize(network, optdata::Dict{TDK, TDV}, ome_optimizer::CodeOptimizer, verbose::Bool) where{TDK, TDV}
-
     try
         @assert TDV <: Number
     catch
@@ -41,30 +40,31 @@ function optimize(network, optdata::Dict{TDK, TDV}, ome_optimizer::CodeOptimizer
     end
 
     # transform the network as EinCode
-    code = network2eincode(network)
+    code, size_dict = network2eincode(network, optdata)
     # optimize the contraction order using OMEinsumContractionOrders, which gives a NestedEinsum
-    optcode = optimize_code(code, optdata, ome_optimizer)
+    optcode = optimize_code(code, size_dict, ome_optimizer)
 
     # transform the optimized contraction order back to the network
     optimaltree = eincode2contractiontree(optcode)
 
-    # calculate the size of maximum tensor during the contraction
-    cc = OMEinsumContractionOrders.contraction_complexity(optcode, optdata)
-    space_complexity = 2.0^cc.sc
-
-    return optimaltree, space_complexity
+    # calculate the complexity of the contraction
+    cc = OMEinsumContractionOrders.contraction_complexity(optcode, size_dict)
+    return optimaltree, 2.0^(cc.tc)
 end
 
-function network2eincode(network)
+function network2eincode(network, optdata)
     indices = unique(vcat(network...))
-    open_edges = empty(network[1])
+    new_indices = Dict([i => j for (j, i) in enumerate(indices)])
+    new_network = [Int[new_indices[i] for i in t] for t in network]
+    open_edges = Int[]
     # if a indices appear only once, it is an open index
     for i in indices
         if sum([i in t for t in network]) == 1
-            push!(open_edges, i)
+            push!(open_edges, new_indices[i])
         end
     end
-    return EinCode(network, open_edges)
+    size_dict = Dict([new_indices[i] => optdata[i] for i in keys(optdata)])
+    return EinCode(new_network, open_edges), size_dict
 end
 
 function eincode2contractiontree(eincode::NestedEinsum)
diff --git a/test/omeinsumcontractionordres.jl b/test/omeinsumcontractionordres.jl
index 76ea819f..7ef0fdb1 100644
--- a/test/omeinsumcontractionordres.jl
+++ b/test/omeinsumcontractionordres.jl
@@ -35,7 +35,7 @@ using OMEinsumContractionOrders
         D5[a, b, c, d] := A[a, e, c, f] * B[g, d, e] * C[g, f, b]
     end
 
-    @tensor opt = (a => 5, b => 5, c => 5, d => 5, e => 5, f => 5, g => 5) opt_algorithm = ExactTreeWidth begin
+    @tensor opt = (a => 5, b => 5, c => 5, d => 5, e => 5, f => 5, g => 5) opt_algorithm = ExactTreewidth begin
         D6[a, b, c, d] := A[a, e, c, f] * B[g, d, e] * C[g, f, b]
     end
 
@@ -73,7 +73,7 @@ using OMEinsumContractionOrders
         s5[] := A[i, k] * B[i, j, l] * C[j, m] * D[k, n] * E[n, l, o] * F[o, m]
     end
 
-    @tensor opt = (i => 2, j => 2, k => 2, l => 2, m => 2, n => 2, o => 2) opt_algorithm = ExactTreeWidth() begin
+    @tensor opt = (i => 2, j => 2, k => 2, l => 2, m => 2, n => 2, o => 2) opt_algorithm = ExactTreewidth begin
         s6[] := A[i, k] * B[i, j, l] * C[j, m] * D[k, n] * E[n, l, o] * F[o, m]
     end
 
@@ -104,7 +104,7 @@ using OMEinsumContractionOrders
         D5[m] := A[i, j, k] * B[j, k, l] * C[i, l, m] + α * A[i, j, k] * B[j, k, l] * C[i, l, m]
     end
 
-    @tensor opt = (i => 5, j => 5, k => 5, l => 5, m => 5) opt_algorithm = ExactTreeWidth begin
+    @tensor opt = (i => 5, j => 5, k => 5, l => 5, m => 5) opt_algorithm = ExactTreewidth begin
         D6[m] := A[i, j, k] * B[j, k, l] * C[i, l, m] + α * A[i, j, k] * B[j, k, l] * C[i, l, m]
     end
 
diff --git a/test/runtests.jl b/test/runtests.jl
index 3b601729..abef21f0 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -8,36 +8,36 @@ using TensorOperations: IndexError
 using TensorOperations: BaseCopy, BaseView, StridedNative, StridedBLAS
 using TensorOperations: DefaultAllocator, ManualAllocator
 
-# @testset "tensoropt" verbose = true begin
-#     include("tensoropt.jl")
-# end
-# @testset "auxiliary" verbose = true begin
-#     include("auxiliary.jl")
-# end
-# @testset "macro keywords" verbose = true begin
-#     include("macro_kwargs.jl")
-# end
-# @testset "method syntax" verbose = true begin
-#     include("methods.jl")
-# end
-# @testset "macro with index notation" verbose = true begin
-#     include("tensor.jl")
-# end
-# @testset "ad" verbose = false begin
-#     include("ad.jl")
-# end
+@testset "tensoropt" verbose = true begin
+    include("tensoropt.jl")
+end
+@testset "auxiliary" verbose = true begin
+    include("auxiliary.jl")
+end
+@testset "macro keywords" verbose = true begin
+    include("macro_kwargs.jl")
+end
+@testset "method syntax" verbose = true begin
+    include("methods.jl")
+end
+@testset "macro with index notation" verbose = true begin
+    include("tensor.jl")
+end
+@testset "ad" verbose = false begin
+    include("ad.jl")
+end
 
-# # note: cuTENSOR should not be loaded before this point
-# # as there is a test which requires it to be loaded after
-# @testset "cuTENSOR extension" verbose = true begin
-#     include("cutensor.jl")
-# end
+# note: cuTENSOR should not be loaded before this point
+# as there is a test which requires it to be loaded after
+@testset "cuTENSOR extension" verbose = true begin
+    include("cutensor.jl")
+end
 
-# # note: Bumper should not be loaded before this point
-# # as there is a test which requires it to be loaded after
-# @testset "Bumper extension" verbose = true begin
-#     include("butensor.jl")
-# end
+# note: Bumper should not be loaded before this point
+# as there is a test which requires it to be loaded after
+@testset "Bumper extension" verbose = true begin
+    include("butensor.jl")
+end
 
 # note: OMEinsumContractionOrders should not be loaded before this point
 # as there is a test which requires it to be loaded after
@@ -45,11 +45,11 @@ using TensorOperations: DefaultAllocator, ManualAllocator
     include("omeinsumcontractionordres.jl")
 end
 
-# @testset "Polynomials" begin
-#     include("polynomials.jl")
-# end
+@testset "Polynomials" begin
+    include("polynomials.jl")
+end
 
-# @testset "Aqua" verbose = true begin
-#     using Aqua
-#     Aqua.test_all(TensorOperations)
-# end
+@testset "Aqua" verbose = true begin
+    using Aqua
+    Aqua.test_all(TensorOperations)
+end

From 3846242a7418fbad0e07c9580b531fe84c796cad Mon Sep 17 00:00:00 2001
From: ArrogantGao <gaoxuanzhao@gmail.com>
Date: Fri, 9 Aug 2024 10:58:35 +0800
Subject: [PATCH 08/15] using kahypar_auto to avoid manully selecting sc_target

---
 ...rOperationsOMEinsumContractionOrdersExt.jl | 35 ++++++++++++++++---
 1 file changed, 31 insertions(+), 4 deletions(-)

diff --git a/ext/TensorOperationsOMEinsumContractionOrdersExt.jl b/ext/TensorOperationsOMEinsumContractionOrdersExt.jl
index 437f1545..8a456dd8 100644
--- a/ext/TensorOperationsOMEinsumContractionOrdersExt.jl
+++ b/ext/TensorOperationsOMEinsumContractionOrdersExt.jl
@@ -4,7 +4,7 @@ using TensorOperations
 using TensorOperations: TensorOperations as TO
 using TensorOperations: TreeOptimizer
 using OMEinsumContractionOrders
-using OMEinsumContractionOrders: EinCode, NestedEinsum, SlicedEinsum, isleaf
+using OMEinsumContractionOrders: EinCode, NestedEinsum, SlicedEinsum, isleaf, optimize_kahypar_auto
 
 function TO.optimaltree(network, optdata::Dict{TDK, TDV}, ::TreeOptimizer{:GreedyMethod}, verbose::Bool) where{TDK, TDV}
     ome_optimizer = GreedyMethod()
@@ -12,9 +12,8 @@ function TO.optimaltree(network, optdata::Dict{TDK, TDV}, ::TreeOptimizer{:Greed
 end
 
 function TO.optimaltree(network, optdata::Dict{TDK, TDV}, ::TreeOptimizer{:KaHyParBipartite}, verbose::Bool) where{TDK, TDV}
-    # sc_target and max_group_size are simply set as 10 for now
-    ome_optimizer = KaHyParBipartite(; sc_target=10, max_group_size=10)
-    return optimize(network, optdata, ome_optimizer, verbose)
+
+    return optimize_kahypar(network, optdata, verbose)
 end
 
 function TO.optimaltree(network, optdata::Dict{TDK, TDV}, ::TreeOptimizer{:TreeSA}, verbose::Bool) where{TDK, TDV}
@@ -49,6 +48,34 @@ function optimize(network, optdata::Dict{TDK, TDV}, ome_optimizer::CodeOptimizer
 
     # calculate the complexity of the contraction
     cc = OMEinsumContractionOrders.contraction_complexity(optcode, size_dict)
+    if verbose
+        println("Optimal contraction tree: ", optimaltree)
+        println(cc)
+    end
+    return optimaltree, 2.0^(cc.tc)
+end
+
+function optimize_kahypar(network, optdata::Dict{TDK, TDV}, verbose::Bool) where{TDK, TDV}
+    try
+        @assert TDV <: Number
+    catch
+        throw(ArgumentError("The values of the optdata dictionary must be of type Number"))
+    end
+
+    # transform the network as EinCode
+    code, size_dict = network2eincode(network, optdata)
+    # optimize the contraction order using OMEinsumContractionOrders, which gives a NestedEinsum
+    optcode = optimize_kahypar_auto(code, size_dict)
+
+    # transform the optimized contraction order back to the network
+    optimaltree = eincode2contractiontree(optcode)
+
+    # calculate the complexity of the contraction
+    cc = OMEinsumContractionOrders.contraction_complexity(optcode, size_dict)
+    if verbose
+        println("Optimal contraction tree: ", optimaltree)
+        println(cc)
+    end
     return optimaltree, 2.0^(cc.tc)
 end
 

From 70301a73d0684cca6bf4c6d9bc7c51a301ea1905 Mon Sep 17 00:00:00 2001
From: ArrogantGao <gaoxuanzhao@gmail.com>
Date: Wed, 14 Aug 2024 16:44:14 +0800
Subject: [PATCH 09/15] update project.toml

---
 Project.toml | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/Project.toml b/Project.toml
index e8f9f6cf..0abbaffe 100644
--- a/Project.toml
+++ b/Project.toml
@@ -64,5 +64,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1"
 
 [targets]
-test = ["Test", "Random", "DynamicPolynomials", "ChainRulesTestUtils", "CUDA", "cuTENSOR", "Aqua", "Logging", "Bumper", "OMEinsumContractionOrders"]
-test = ["Test", "Random", "DynamicPolynomials", "ChainRulesTestUtils", "CUDA", "cuTENSOR", "Aqua", "Logging", "Bumper"]
+test = ["Test", "Random", "DynamicPolynomials", "ChainRulesTestUtils", "CUDA", "cuTENSOR", "Aqua", "Logging", "Bumper", "OMEinsumContractionOrders"]
\ No newline at end of file

From 61d1cce30e831f98dbe06c27bd3edfcde626014a Mon Sep 17 00:00:00 2001
From: ArrogantGao <gaoxuanzhao@gmail.com>
Date: Wed, 14 Aug 2024 23:47:24 +0800
Subject: [PATCH 10/15] add optimizer to ncon

---
 .github/workflows/ci.yml                      |  2 +-
 ...rOperationsOMEinsumContractionOrdersExt.jl | 18 +++---
 src/implementation/ncon.jl                    | 35 +++++++++++-
 src/indexnotation/optimaltree.jl              |  9 +--
 src/indexnotation/tensormacros.jl             | 28 +++++----
 test/macro_kwargs.jl                          |  4 +-
 test/omeinsumcontractionordres.jl             | 57 +++++++++++++++----
 7 files changed, 111 insertions(+), 42 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 84322407..47f5ad31 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -21,7 +21,7 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.8' # lowest supported version
+          - '1.9' # lowest supported version
           - '1' # automatically expands to the latest stable 1.x release of Julia
         os:
           - ubuntu-latest
diff --git a/ext/TensorOperationsOMEinsumContractionOrdersExt.jl b/ext/TensorOperationsOMEinsumContractionOrdersExt.jl
index 8a456dd8..5102783b 100644
--- a/ext/TensorOperationsOMEinsumContractionOrdersExt.jl
+++ b/ext/TensorOperationsOMEinsumContractionOrdersExt.jl
@@ -7,36 +7,36 @@ using OMEinsumContractionOrders
 using OMEinsumContractionOrders: EinCode, NestedEinsum, SlicedEinsum, isleaf, optimize_kahypar_auto
 
 function TO.optimaltree(network, optdata::Dict{TDK, TDV}, ::TreeOptimizer{:GreedyMethod}, verbose::Bool) where{TDK, TDV}
+    @debug "Using optimizer GreedyMethod from OMEinsumContractionOrders"
     ome_optimizer = GreedyMethod()
     return optimize(network, optdata, ome_optimizer, verbose)
 end
 
 function TO.optimaltree(network, optdata::Dict{TDK, TDV}, ::TreeOptimizer{:KaHyParBipartite}, verbose::Bool) where{TDK, TDV}
-
+    @debug "Using optimizer KaHyParBipartite from OMEinsumContractionOrders"
     return optimize_kahypar(network, optdata, verbose)
 end
 
 function TO.optimaltree(network, optdata::Dict{TDK, TDV}, ::TreeOptimizer{:TreeSA}, verbose::Bool) where{TDK, TDV}
+    @debug "Using optimizer TreeSA from OMEinsumContractionOrders"
     ome_optimizer = TreeSA()
     return optimize(network, optdata, ome_optimizer, verbose)
 end
 
 function TO.optimaltree(network, optdata::Dict{TDK, TDV}, ::TreeOptimizer{:SABipartite}, verbose::Bool) where{TDK, TDV}
+    @debug "Using optimizer SABipartite from OMEinsumContractionOrders"
     ome_optimizer = SABipartite()
     return optimize(network, optdata, ome_optimizer, verbose)
 end
 
 function TO.optimaltree(network, optdata::Dict{TDK, TDV}, ::TreeOptimizer{:ExactTreewidth}, verbose::Bool) where{TDK, TDV}
+    @debug "Using optimizer ExactTreewidth from OMEinsumContractionOrders"
     ome_optimizer = ExactTreewidth()
     return optimize(network, optdata, ome_optimizer, verbose)
 end
 
 function optimize(network, optdata::Dict{TDK, TDV}, ome_optimizer::CodeOptimizer, verbose::Bool) where{TDK, TDV}
-    try
-        @assert TDV <: Number
-    catch
-        throw(ArgumentError("The values of the optdata dictionary must be of type Number"))
-    end
+    @assert TDV <: Number "The values of `optdata` dictionary must be of `<:Number`"
 
     # transform the network as EinCode
     code, size_dict = network2eincode(network, optdata)
@@ -56,11 +56,7 @@ function optimize(network, optdata::Dict{TDK, TDV}, ome_optimizer::CodeOptimizer
 end
 
 function optimize_kahypar(network, optdata::Dict{TDK, TDV}, verbose::Bool) where{TDK, TDV}
-    try
-        @assert TDV <: Number
-    catch
-        throw(ArgumentError("The values of the optdata dictionary must be of type Number"))
-    end
+    @assert TDV <: Number "The values of `optdata` dictionary must be of `<:Number`"
 
     # transform the network as EinCode
     code, size_dict = network2eincode(network, optdata)
diff --git a/src/implementation/ncon.jl b/src/implementation/ncon.jl
index 7dec5532..4b70458a 100644
--- a/src/implementation/ncon.jl
+++ b/src/implementation/ncon.jl
@@ -1,5 +1,6 @@
 """
     ncon(tensorlist, indexlist, [conjlist, sym]; order = ..., output = ..., backend = ..., allocator = ...)
+    ncon(tensorlist, indexlist, optimizer, conjlist; output=..., kwargs...)
 
 Contract the tensors in `tensorlist` (of type `Vector` or `Tuple`) according to the network
 as specified by `indexlist`. Here, `indexlist` is a list (i.e. a `Vector` or `Tuple`) with
@@ -20,6 +21,9 @@ over are labelled by increasing integers, i.e. first the contraction correspondi
 (negative, so increasing in absolute value) index labels. The keyword arguments `order` and
 `output` allow to change these defaults.
 
+Another way to get the contraction order is to use the TreeOptimizer, by passing the `optimizer` (which a Symbol) instead of the `order` keyword argument. The `optimizer` can be `:ExhaustiveSearch`.
+With the extension `OMEinsumContractionOrders`, the `optimizer` can be one of the following: `:GreedyMethod`, `:TreeSA`, `:KaHyParBipartite`, `:SABipartite`, `:ExactTreewidth`.
+
 See also the macro version [`@ncon`](@ref).
 """
 function ncon(tensors, network,
@@ -40,10 +44,39 @@ function ncon(tensors, network,
 
     (tensors, network) = resolve_traces(tensors, network)
     tree = order === nothing ? ncontree(network) : indexordertree(network, order)
+    return ncon(tensors, network, conjlist, tree, output′; kwargs...)
+end
+function ncon(tensors, network, optimizer::T, conjlist=fill(false, length(tensors)); output=nothing, kwargs...) where{T <: Symbol}
+    length(tensors) == length(network) == length(conjlist) ||
+        throw(ArgumentError("number of tensors and of index lists should be the same"))
+    isnconstyle(network) || throw(ArgumentError("invalid NCON network: $network"))
+    output′ = nconoutput(network, output)
+
+    if length(tensors) == 1
+        if length(output′) == length(network[1])
+            return tensorcopy(output′, tensors[1], network[1], conjlist[1]; kwargs...)
+        else
+            return tensortrace(output′, tensors[1], network[1], conjlist[1]; kwargs...)
+        end
+    end
+
+    (tensors, network) = resolve_traces(tensors, network)
+
+    optdata = Dict{Any, Number}()
+    for (i, ids) in enumerate(network)
+        for (j, id) in enumerate(ids)
+            optdata[id] = size(tensors[i], j)
+        end
+    end
+
+    tree = optimaltree(network, optdata, TreeOptimizer{optimizer}(), false)[1]
+    return ncon(tensors, network, conjlist, tree, output′; kwargs...)
+end
 
+function ncon(tensors, network, conjlist, tree, output; kwargs...)
     A, IA, conjA = contracttree(tensors, network, conjlist, tree[1]; kwargs...)
     B, IB, conjB = contracttree(tensors, network, conjlist, tree[2]; kwargs...)
-    IC = tuple(output′...)
+    IC = tuple(output...)
     C = tensorcontract(IC, A, IA, conjA, B, IB, conjB; kwargs...)
     allocator = haskey(kwargs, :allocator) ? kwargs[:allocator] : DefaultAllocator()
     tree[1] isa Int || tensorfree!(A, allocator)
diff --git a/src/indexnotation/optimaltree.jl b/src/indexnotation/optimaltree.jl
index ac4048c6..4166f0db 100644
--- a/src/indexnotation/optimaltree.jl
+++ b/src/indexnotation/optimaltree.jl
@@ -1,15 +1,16 @@
 struct TreeOptimizer{T} end # T is a Symbol for the algorithm
 
-function optimaltree(network, optdata::Dict; optimizer::TreeOptimizer = TreeOptimizer{:NCon}(), verbose::Bool=false)
-    @debug "Using optimizer $(typeof(optimizer))"
+function optimaltree(network, optdata::Dict;
+                     optimizer::TreeOptimizer{T}=TreeOptimizer{:ExhaustiveSearch}(), verbose::Bool=false) where{T}
     return optimaltree(network, optdata, optimizer, verbose)
 end
 
-function optimaltree(network, optdata::Dict, ::TreeOptimizer{T}, verbose::Bool) where{T}
+function optimaltree(network, optdata::Dict, ::TreeOptimizer{T}, verbose::Bool) where {T}
     throw(ArgumentError("Unknown optimizer: $T. Hint: may need to load extensions, e.g. `using OMEinsumContractionOrders`"))
 end
 
-function optimaltree(network, optdata::Dict, ::TreeOptimizer{:NCon}, verbose::Bool)
+function optimaltree(network, optdata::Dict, ::TreeOptimizer{:ExhaustiveSearch}, verbose::Bool)
+    @debug "Using optimizer ExhaustiveSearch"
     numtensors = length(network)
     allindices = unique(vcat(network...))
     numindices = length(allindices)
diff --git a/src/indexnotation/tensormacros.jl b/src/indexnotation/tensormacros.jl
index 4581e2dd..c3ee864e 100644
--- a/src/indexnotation/tensormacros.jl
+++ b/src/indexnotation/tensormacros.jl
@@ -71,7 +71,8 @@ function tensorparser(tensorexpr, kwargs...)
         end
     end
     # now handle the remaining keyword arguments
-    optimizer = TreeOptimizer{:NCon}() # the default optimizer
+    optimizer = TreeOptimizer{:ExhaustiveSearch}() # the default optimizer implemented in TensorOperations.jl
+    optval = nothing
     for (name, val) in kwargs
         if name == :order
             isexpr(val, :tuple) ||
@@ -86,25 +87,30 @@ function tensorparser(tensorexpr, kwargs...)
             val in (:warn, :cache) ||
                 throw(ArgumentError("Invalid use of `costcheck`, should be `costcheck=warn` or `costcheck=cache`"))
             parser.contractioncostcheck = val
+        elseif name == :opt
+            optval = val
         elseif name == :opt_algorithm
             if val isa Symbol
                 optimizer = TreeOptimizer{val}()
             else
-                throw(ArgumentError("Invalid use of `opt_algorithm`, should be `opt_algorithm=NCon` or `opt_algorithm=NameOfAlgorithm`"))
-            end
-        elseif name == :opt
-            if val isa Bool && val
-                optdict = optdata(tensorexpr)
-            elseif val isa Expr
-                optdict = optdata(val, tensorexpr)
-            else
-                throw(ArgumentError("Invalid use of `opt`, should be `opt=true` or `opt=OptExpr`"))
+                throw(ArgumentError("Invalid use of `opt_algorithm`, should be `opt_algorithm=ExhaustiveSearch` or `opt_algorithm=NameOfAlgorithm`"))
             end
-            parser.contractiontreebuilder = network -> optimaltree(network, optdict, optimizer = optimizer)[1]
         elseif !(name == :backend || name == :allocator) # these two have been handled
             throw(ArgumentError("Unknown keyword argument `name`."))
         end
     end
+    # construct the contraction tree builder after all keyword arguments have been processed
+    if !isnothing(optval)
+        if optval isa Bool && optval
+            optdict = optdata(tensorexpr)
+        elseif optval isa Expr
+            optdict = optdata(optval, tensorexpr)
+        else
+            throw(ArgumentError("Invalid use of `opt`, should be `opt=true` or `opt=OptExpr`"))
+        end
+        parser.contractiontreebuilder = network -> optimaltree(network, optdict;
+                                                            optimizer=optimizer)[1]
+    end
     return parser
 end
 
diff --git a/test/macro_kwargs.jl b/test/macro_kwargs.jl
index 64826d9a..4ecddd0f 100644
--- a/test/macro_kwargs.jl
+++ b/test/macro_kwargs.jl
@@ -116,9 +116,9 @@ end
         D1[a, b, c, d] := A[a, e, c, f] * B[g, d, e] * C[g, f, b]
     end
 
-    @tensor opt = true opt_algorithm = NCon begin
+    @tensor opt = true opt_algorithm = ExhaustiveSearch begin
         D2[a, b, c, d] := A[a, e, c, f] * B[g, d, e] * C[g, f, b]
     end
 
     @test D1 ≈ D2
-end
\ No newline at end of file
+end
diff --git a/test/omeinsumcontractionordres.jl b/test/omeinsumcontractionordres.jl
index 7ef0fdb1..ec73610d 100644
--- a/test/omeinsumcontractionordres.jl
+++ b/test/omeinsumcontractionordres.jl
@@ -1,15 +1,18 @@
 @testset "@tensor dependency check" begin
+    A = rand(2, 2)
+    B = rand(2, 2)
+    C = rand(2, 2)
     @test_throws ArgumentError begin
-        A = rand(2, 2)
-        B = rand(2, 2)
-        C = rand(2, 2)
-        ex = :(@tensor opt=(i=>2, j=>2, k=>2) opt_algorithm=GreedyMethod S[] := A[i, j] * B[j, k] * C[i, k])
+        ex = :(@tensor opt = (i => 2, j => 2, k => 2) opt_algorithm = GreedyMethod S[] := A[i,j] * B[j,k] * C[i, k])
         macroexpand(Main, ex)
     end
 end
 
 using OMEinsumContractionOrders
 
+# open the debug mode to check the optimization algorithms used
+ENV["JULIA_DEBUG"] = "TensorOperationsOMEinsumContractionOrdersExt"
+
 @testset "OMEinsumContractionOrders optimization algorithms" begin
     A = randn(5, 5, 5, 5)
     B = randn(5, 5, 5)
@@ -43,8 +46,12 @@ using OMEinsumContractionOrders
         D7[1, 2, 3, 4] := A[1, 5, 3, 6] * B[7, 4, 5] * C[7, 6, 2]
     end
 
-    @test D1 ≈ D2 ≈ D3 ≈ D4 ≈ D5 ≈ D6 ≈ D7
+    # check the case that opt_algorithm is before the opt
+    @tensor opt_algorithm = GreedyMethod opt = (a => 5, b => 5, c => 5, d => 5, e => 5, f => 5, g => 5) begin
+        D8[a, b, c, d] := A[a, e, c, f] * B[g, d, e] * C[g, f, b]
+    end
 
+    @test D1 ≈ D2 ≈ D3 ≈ D4 ≈ D5 ≈ D6 ≈ D7 ≈ D8
 
     A = rand(2, 2)
     B = rand(2, 2, 2)
@@ -52,7 +59,7 @@ using OMEinsumContractionOrders
     D = rand(2, 2)
     E = rand(2, 2, 2)
     F = rand(2, 2)
-    
+
     @tensor opt = true begin
         s1[] := A[i, k] * B[i, j, l] * C[j, m] * D[k, n] * E[n, l, o] * F[o, m]
     end
@@ -85,28 +92,54 @@ using OMEinsumContractionOrders
     α = randn()
 
     @tensor opt = true begin
-        D1[m] := A[i, j, k] * B[j, k, l] * C[i, l, m] + α * A[i, j, k] * B[j, k, l] * C[i, l, m]
+        D1[m] := A[i, j, k] * B[j, k, l] * C[i, l, m] +
+                 α * A[i, j, k] * B[j, k, l] * C[i, l, m]
     end
 
     @tensor opt = (i => 5, j => 5, k => 5, l => 5, m => 5) opt_algorithm = GreedyMethod begin
-        D2[m] := A[i, j, k] * B[j, k, l] * C[i, l, m] + α * A[i, j, k] * B[j, k, l] * C[i, l, m]
+        D2[m] := A[i, j, k] * B[j, k, l] * C[i, l, m] +
+                 α * A[i, j, k] * B[j, k, l] * C[i, l, m]
     end
 
     @tensor opt = (i => 5, j => 5, k => 5, l => 5, m => 5) opt_algorithm = TreeSA begin
-        D3[m] := A[i, j, k] * B[j, k, l] * C[i, l, m] + α * A[i, j, k] * B[j, k, l] * C[i, l, m]
+        D3[m] := A[i, j, k] * B[j, k, l] * C[i, l, m] +
+                 α * A[i, j, k] * B[j, k, l] * C[i, l, m]
     end
 
     @tensor opt = (i => 5, j => 5, k => 5, l => 5, m => 5) opt_algorithm = KaHyParBipartite begin
-        D4[m] := A[i, j, k] * B[j, k, l] * C[i, l, m] + α * A[i, j, k] * B[j, k, l] * C[i, l, m]
+        D4[m] := A[i, j, k] * B[j, k, l] * C[i, l, m] +
+                 α * A[i, j, k] * B[j, k, l] * C[i, l, m]
     end
 
     @tensor opt = (i => 5, j => 5, k => 5, l => 5, m => 5) opt_algorithm = SABipartite begin
-        D5[m] := A[i, j, k] * B[j, k, l] * C[i, l, m] + α * A[i, j, k] * B[j, k, l] * C[i, l, m]
+        D5[m] := A[i, j, k] * B[j, k, l] * C[i, l, m] +
+                 α * A[i, j, k] * B[j, k, l] * C[i, l, m]
     end
 
     @tensor opt = (i => 5, j => 5, k => 5, l => 5, m => 5) opt_algorithm = ExactTreewidth begin
-        D6[m] := A[i, j, k] * B[j, k, l] * C[i, l, m] + α * A[i, j, k] * B[j, k, l] * C[i, l, m]
+        D6[m] := A[i, j, k] * B[j, k, l] * C[i, l, m] +
+                 α * A[i, j, k] * B[j, k, l] * C[i, l, m]
     end
 
     @test D1 ≈ D2 ≈ D3 ≈ D4 ≈ D5 ≈ D6
+end
+
+@testset "ncon with OMEinsumContractionOrders" begin
+    A = randn(5, 5, 5, 5)
+    B = randn(5, 5, 5)
+    C = randn(5, 5, 5)
+
+    @tensor begin
+        D1[a, b, c, d] := A[a, e, c, f] * B[g, d, e] * C[g, f, b]
+    end
+
+    D2 = ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]])
+    D3 = ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]], :ExhaustiveSearch)
+    D4 = ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]], :GreedyMethod)
+    D5 = ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]], :KaHyParBipartite)
+    D6 = ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]], :TreeSA)
+    D7 = ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]], :SABipartite)
+    D8 = ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]], :ExactTreewidth)
+
+    @test D1 ≈ D2 ≈ D3 ≈ D4 ≈ D5 ≈ D6 ≈ D7 ≈ D8
 end
\ No newline at end of file

From f55c1b22a5726fa1ec31f20b99db90e4a97561ec Mon Sep 17 00:00:00 2001
From: ArrogantGao <gaoxuanzhao@gmail.com>
Date: Thu, 15 Aug 2024 00:04:19 +0800
Subject: [PATCH 11/15] formatted by JuliaFormatter

---
 ...rOperationsOMEinsumContractionOrdersExt.jl | 25 ++++++++++++-------
 src/implementation/ncon.jl                    | 11 +++++---
 src/indexnotation/optimaltree.jl              |  6 +++--
 src/indexnotation/tensormacros.jl             |  2 +-
 test/omeinsumcontractionordres.jl             | 12 ++++++---
 5 files changed, 37 insertions(+), 19 deletions(-)

diff --git a/ext/TensorOperationsOMEinsumContractionOrdersExt.jl b/ext/TensorOperationsOMEinsumContractionOrdersExt.jl
index 5102783b..1e846d99 100644
--- a/ext/TensorOperationsOMEinsumContractionOrdersExt.jl
+++ b/ext/TensorOperationsOMEinsumContractionOrdersExt.jl
@@ -4,38 +4,45 @@ using TensorOperations
 using TensorOperations: TensorOperations as TO
 using TensorOperations: TreeOptimizer
 using OMEinsumContractionOrders
-using OMEinsumContractionOrders: EinCode, NestedEinsum, SlicedEinsum, isleaf, optimize_kahypar_auto
+using OMEinsumContractionOrders: EinCode, NestedEinsum, SlicedEinsum, isleaf,
+                                 optimize_kahypar_auto
 
-function TO.optimaltree(network, optdata::Dict{TDK, TDV}, ::TreeOptimizer{:GreedyMethod}, verbose::Bool) where{TDK, TDV}
+function TO.optimaltree(network, optdata::Dict{TDK,TDV}, ::TreeOptimizer{:GreedyMethod},
+                        verbose::Bool) where {TDK,TDV}
     @debug "Using optimizer GreedyMethod from OMEinsumContractionOrders"
     ome_optimizer = GreedyMethod()
     return optimize(network, optdata, ome_optimizer, verbose)
 end
 
-function TO.optimaltree(network, optdata::Dict{TDK, TDV}, ::TreeOptimizer{:KaHyParBipartite}, verbose::Bool) where{TDK, TDV}
+function TO.optimaltree(network, optdata::Dict{TDK,TDV}, ::TreeOptimizer{:KaHyParBipartite},
+                        verbose::Bool) where {TDK,TDV}
     @debug "Using optimizer KaHyParBipartite from OMEinsumContractionOrders"
     return optimize_kahypar(network, optdata, verbose)
 end
 
-function TO.optimaltree(network, optdata::Dict{TDK, TDV}, ::TreeOptimizer{:TreeSA}, verbose::Bool) where{TDK, TDV}
+function TO.optimaltree(network, optdata::Dict{TDK,TDV}, ::TreeOptimizer{:TreeSA},
+                        verbose::Bool) where {TDK,TDV}
     @debug "Using optimizer TreeSA from OMEinsumContractionOrders"
     ome_optimizer = TreeSA()
     return optimize(network, optdata, ome_optimizer, verbose)
 end
 
-function TO.optimaltree(network, optdata::Dict{TDK, TDV}, ::TreeOptimizer{:SABipartite}, verbose::Bool) where{TDK, TDV}
+function TO.optimaltree(network, optdata::Dict{TDK,TDV}, ::TreeOptimizer{:SABipartite},
+                        verbose::Bool) where {TDK,TDV}
     @debug "Using optimizer SABipartite from OMEinsumContractionOrders"
     ome_optimizer = SABipartite()
     return optimize(network, optdata, ome_optimizer, verbose)
 end
 
-function TO.optimaltree(network, optdata::Dict{TDK, TDV}, ::TreeOptimizer{:ExactTreewidth}, verbose::Bool) where{TDK, TDV}
+function TO.optimaltree(network, optdata::Dict{TDK,TDV}, ::TreeOptimizer{:ExactTreewidth},
+                        verbose::Bool) where {TDK,TDV}
     @debug "Using optimizer ExactTreewidth from OMEinsumContractionOrders"
     ome_optimizer = ExactTreewidth()
     return optimize(network, optdata, ome_optimizer, verbose)
 end
 
-function optimize(network, optdata::Dict{TDK, TDV}, ome_optimizer::CodeOptimizer, verbose::Bool) where{TDK, TDV}
+function optimize(network, optdata::Dict{TDK,TDV}, ome_optimizer::CodeOptimizer,
+                  verbose::Bool) where {TDK,TDV}
     @assert TDV <: Number "The values of `optdata` dictionary must be of `<:Number`"
 
     # transform the network as EinCode
@@ -55,7 +62,7 @@ function optimize(network, optdata::Dict{TDK, TDV}, ome_optimizer::CodeOptimizer
     return optimaltree, 2.0^(cc.tc)
 end
 
-function optimize_kahypar(network, optdata::Dict{TDK, TDV}, verbose::Bool) where{TDK, TDV}
+function optimize_kahypar(network, optdata::Dict{TDK,TDV}, verbose::Bool) where {TDK,TDV}
     @assert TDV <: Number "The values of `optdata` dictionary must be of `<:Number`"
 
     # transform the network as EinCode
@@ -103,4 +110,4 @@ function eincode2contractiontree(eincode::SlicedEinsum)
     return eincode2contractiontree(eincode.eins)
 end
 
-end
\ No newline at end of file
+end
diff --git a/src/implementation/ncon.jl b/src/implementation/ncon.jl
index 4b70458a..502d44e5 100644
--- a/src/implementation/ncon.jl
+++ b/src/implementation/ncon.jl
@@ -21,8 +21,10 @@ over are labelled by increasing integers, i.e. first the contraction correspondi
 (negative, so increasing in absolute value) index labels. The keyword arguments `order` and
 `output` allow to change these defaults.
 
-Another way to get the contraction order is to use the TreeOptimizer, by passing the `optimizer` (which a Symbol) instead of the `order` keyword argument. The `optimizer` can be `:ExhaustiveSearch`.
-With the extension `OMEinsumContractionOrders`, the `optimizer` can be one of the following: `:GreedyMethod`, `:TreeSA`, `:KaHyParBipartite`, `:SABipartite`, `:ExactTreewidth`.
+Another way to get the contraction order is to use the TreeOptimizer, by passing the `optimizer` 
+(which a Symbol) instead of the `order` keyword argument. The `optimizer` can be `:ExhaustiveSearch`.
+With the extension `OMEinsumContractionOrders`, the `optimizer` can be one of the following: 
+`:GreedyMethod`, `:TreeSA`, `:KaHyParBipartite`, `:SABipartite`, `:ExactTreewidth`.
 
 See also the macro version [`@ncon`](@ref).
 """
@@ -46,7 +48,8 @@ function ncon(tensors, network,
     tree = order === nothing ? ncontree(network) : indexordertree(network, order)
     return ncon(tensors, network, conjlist, tree, output′; kwargs...)
 end
-function ncon(tensors, network, optimizer::T, conjlist=fill(false, length(tensors)); output=nothing, kwargs...) where{T <: Symbol}
+function ncon(tensors, network, optimizer::T, conjlist=fill(false, length(tensors));
+              output=nothing, kwargs...) where {T<:Symbol}
     length(tensors) == length(network) == length(conjlist) ||
         throw(ArgumentError("number of tensors and of index lists should be the same"))
     isnconstyle(network) || throw(ArgumentError("invalid NCON network: $network"))
@@ -62,7 +65,7 @@ function ncon(tensors, network, optimizer::T, conjlist=fill(false, length(tensor
 
     (tensors, network) = resolve_traces(tensors, network)
 
-    optdata = Dict{Any, Number}()
+    optdata = Dict{Any,Number}()
     for (i, ids) in enumerate(network)
         for (j, id) in enumerate(ids)
             optdata[id] = size(tensors[i], j)
diff --git a/src/indexnotation/optimaltree.jl b/src/indexnotation/optimaltree.jl
index 4166f0db..124a0c7d 100644
--- a/src/indexnotation/optimaltree.jl
+++ b/src/indexnotation/optimaltree.jl
@@ -1,7 +1,8 @@
 struct TreeOptimizer{T} end # T is a Symbol for the algorithm
 
 function optimaltree(network, optdata::Dict;
-                     optimizer::TreeOptimizer{T}=TreeOptimizer{:ExhaustiveSearch}(), verbose::Bool=false) where{T}
+                     optimizer::TreeOptimizer{T}=TreeOptimizer{:ExhaustiveSearch}(),
+                     verbose::Bool=false) where {T}
     return optimaltree(network, optdata, optimizer, verbose)
 end
 
@@ -9,7 +10,8 @@ function optimaltree(network, optdata::Dict, ::TreeOptimizer{T}, verbose::Bool)
     throw(ArgumentError("Unknown optimizer: $T. Hint: may need to load extensions, e.g. `using OMEinsumContractionOrders`"))
 end
 
-function optimaltree(network, optdata::Dict, ::TreeOptimizer{:ExhaustiveSearch}, verbose::Bool)
+function optimaltree(network, optdata::Dict, ::TreeOptimizer{:ExhaustiveSearch},
+                     verbose::Bool)
     @debug "Using optimizer ExhaustiveSearch"
     numtensors = length(network)
     allindices = unique(vcat(network...))
diff --git a/src/indexnotation/tensormacros.jl b/src/indexnotation/tensormacros.jl
index c3ee864e..7b11913b 100644
--- a/src/indexnotation/tensormacros.jl
+++ b/src/indexnotation/tensormacros.jl
@@ -109,7 +109,7 @@ function tensorparser(tensorexpr, kwargs...)
             throw(ArgumentError("Invalid use of `opt`, should be `opt=true` or `opt=OptExpr`"))
         end
         parser.contractiontreebuilder = network -> optimaltree(network, optdict;
-                                                            optimizer=optimizer)[1]
+                                                               optimizer=optimizer)[1]
     end
     return parser
 end
diff --git a/test/omeinsumcontractionordres.jl b/test/omeinsumcontractionordres.jl
index ec73610d..f8280e38 100644
--- a/test/omeinsumcontractionordres.jl
+++ b/test/omeinsumcontractionordres.jl
@@ -3,7 +3,12 @@
     B = rand(2, 2)
     C = rand(2, 2)
     @test_throws ArgumentError begin
-        ex = :(@tensor opt = (i => 2, j => 2, k => 2) opt_algorithm = GreedyMethod S[] := A[i,j] * B[j,k] * C[i, k])
+        ex = :(@tensor opt = (i => 2, j => 2, k => 2) opt_algorithm = GreedyMethod S[] := A[i,
+                                                                                            j] *
+                                                                                          B[j,
+                                                                                            k] *
+                                                                                          C[i,
+                                                                                            k])
         macroexpand(Main, ex)
     end
 end
@@ -47,7 +52,8 @@ ENV["JULIA_DEBUG"] = "TensorOperationsOMEinsumContractionOrdersExt"
     end
 
     # check the case that opt_algorithm is before the opt
-    @tensor opt_algorithm = GreedyMethod opt = (a => 5, b => 5, c => 5, d => 5, e => 5, f => 5, g => 5) begin
+    @tensor opt_algorithm = GreedyMethod opt = (a => 5, b => 5, c => 5, d => 5, e => 5,
+                                                f => 5, g => 5) begin
         D8[a, b, c, d] := A[a, e, c, f] * B[g, d, e] * C[g, f, b]
     end
 
@@ -142,4 +148,4 @@ end
     D8 = ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]], :ExactTreewidth)
 
     @test D1 ≈ D2 ≈ D3 ≈ D4 ≈ D5 ≈ D6 ≈ D7 ≈ D8
-end
\ No newline at end of file
+end

From 92fe983bd33ac57674e68cc9e870c20d3dac135b Mon Sep 17 00:00:00 2001
From: ArrogantGao <gaoxuanzhao@gmail.com>
Date: Thu, 15 Aug 2024 12:04:00 +0800
Subject: [PATCH 12/15] export TreeOptimizer

---
 .github/workflows/ci.yml                      |  2 +-
 Project.toml                                  |  2 +-
 src/TensorOperations.jl                       |  3 ++
 src/implementation/ncon.jl                    | 47 +++++++++----------
 src/indexnotation/optimaltree.jl              |  6 +++
 ...einsumcontractionordres.jl => omeinsum.jl} | 30 ++++++++----
 test/runtests.jl                              |  8 +++-
 7 files changed, 60 insertions(+), 38 deletions(-)
 rename test/{omeinsumcontractionordres.jl => omeinsum.jl} (77%)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 47f5ad31..84322407 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -21,7 +21,7 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.9' # lowest supported version
+          - '1.8' # lowest supported version
           - '1' # automatically expands to the latest stable 1.x release of Julia
         os:
           - ubuntu-latest
diff --git a/Project.toml b/Project.toml
index 0abbaffe..3934e7b2 100644
--- a/Project.toml
+++ b/Project.toml
@@ -64,4 +64,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 cuTENSOR = "011b41b2-24ef-40a8-b3eb-fa098493e9e1"
 
 [targets]
-test = ["Test", "Random", "DynamicPolynomials", "ChainRulesTestUtils", "CUDA", "cuTENSOR", "Aqua", "Logging", "Bumper", "OMEinsumContractionOrders"]
\ No newline at end of file
+test = ["Test", "Random", "DynamicPolynomials", "ChainRulesTestUtils", "CUDA", "cuTENSOR", "Aqua", "Logging", "Bumper", "OMEinsumContractionOrders"]
diff --git a/src/TensorOperations.jl b/src/TensorOperations.jl
index d3e6ff7d..583f9853 100644
--- a/src/TensorOperations.jl
+++ b/src/TensorOperations.jl
@@ -29,6 +29,9 @@ export IndexTuple, Index2Tuple, linearize
 # export debug functionality
 export checkcontractible, tensorcost
 
+# export optimizer
+export TreeOptimizer, ExhaustiveSearchOptimizer, GreedyMethodOptimizer, KaHyParBipartiteOptimizer, TreeSAOptimizer, SABipartiteOptimizer, ExactTreewidthOptimizer
+
 # Interface and index types
 #---------------------------
 include("indices.jl")
diff --git a/src/implementation/ncon.jl b/src/implementation/ncon.jl
index 502d44e5..42d28f22 100644
--- a/src/implementation/ncon.jl
+++ b/src/implementation/ncon.jl
@@ -1,6 +1,5 @@
 """
-    ncon(tensorlist, indexlist, [conjlist, sym]; order = ..., output = ..., backend = ..., allocator = ...)
-    ncon(tensorlist, indexlist, optimizer, conjlist; output=..., kwargs...)
+    ncon(tensorlist, indexlist, [conjlist, sym]; order = ..., output = ..., optimizer = ..., backend = ..., allocator = ...)
 
 Contract the tensors in `tensorlist` (of type `Vector` or `Tuple`) according to the network
 as specified by `indexlist`. Here, `indexlist` is a list (i.e. a `Vector` or `Tuple`) with
@@ -22,7 +21,7 @@ over are labelled by increasing integers, i.e. first the contraction correspondi
 `output` allow to change these defaults.
 
 Another way to get the contraction order is to use the TreeOptimizer, by passing the `optimizer` 
-(which a Symbol) instead of the `order` keyword argument. The `optimizer` can be `:ExhaustiveSearch`.
+instead of the `order` keyword argument. The `optimizer` can be `:ExhaustiveSearch`.
 With the extension `OMEinsumContractionOrders`, the `optimizer` can be one of the following: 
 `:GreedyMethod`, `:TreeSA`, `:KaHyParBipartite`, `:SABipartite`, `:ExactTreewidth`.
 
@@ -30,7 +29,7 @@ See also the macro version [`@ncon`](@ref).
 """
 function ncon(tensors, network,
               conjlist=fill(false, length(tensors));
-              order=nothing, output=nothing, kwargs...)
+              order=nothing, output=nothing, optimizer=nothing, kwargs...)
     length(tensors) == length(network) == length(conjlist) ||
         throw(ArgumentError("number of tensors and of index lists should be the same"))
     isnconstyle(network) || throw(ArgumentError("invalid NCON network: $network"))
@@ -45,34 +44,30 @@ function ncon(tensors, network,
     end
 
     (tensors, network) = resolve_traces(tensors, network)
-    tree = order === nothing ? ncontree(network) : indexordertree(network, order)
-    return ncon(tensors, network, conjlist, tree, output′; kwargs...)
-end
-function ncon(tensors, network, optimizer::T, conjlist=fill(false, length(tensors));
-              output=nothing, kwargs...) where {T<:Symbol}
-    length(tensors) == length(network) == length(conjlist) ||
-        throw(ArgumentError("number of tensors and of index lists should be the same"))
-    isnconstyle(network) || throw(ArgumentError("invalid NCON network: $network"))
-    output′ = nconoutput(network, output)
 
-    if length(tensors) == 1
-        if length(output′) == length(network[1])
-            return tensorcopy(output′, tensors[1], network[1], conjlist[1]; kwargs...)
+    if isnothing(order)
+        if isnothing(optimizer)
+            # not specifing order and optimizer, tree via ncontree
+            tree=ncontree(network)
         else
-            return tensortrace(output′, tensors[1], network[1], conjlist[1]; kwargs...)
+            # order via optimizer
+            optdata = Dict{Any,Number}()
+            for (i, ids) in enumerate(network)
+                for (j, id) in enumerate(ids)
+                    optdata[id] = tensorstructure(tensors[i], j, conjlist[i])
+                end
+            end
+            tree=optimaltree(network, optdata, optimizer, false)[1]
         end
-    end
-
-    (tensors, network) = resolve_traces(tensors, network)
-
-    optdata = Dict{Any,Number}()
-    for (i, ids) in enumerate(network)
-        for (j, id) in enumerate(ids)
-            optdata[id] = size(tensors[i], j)
+    else
+        if !isnothing(optimizer)
+            throw(ArgumentError("cannot specify both `order` and `optimizer`"))
+        else
+            # with given order, tree via indexordertree
+            tree = indexordertree(network, order)
         end
     end
 
-    tree = optimaltree(network, optdata, TreeOptimizer{optimizer}(), false)[1]
     return ncon(tensors, network, conjlist, tree, output′; kwargs...)
 end
 
diff --git a/src/indexnotation/optimaltree.jl b/src/indexnotation/optimaltree.jl
index 124a0c7d..a3e8ceee 100644
--- a/src/indexnotation/optimaltree.jl
+++ b/src/indexnotation/optimaltree.jl
@@ -1,4 +1,10 @@
 struct TreeOptimizer{T} end # T is a Symbol for the algorithm
+ExhaustiveSearchOptimizer() = TreeOptimizer{:ExhaustiveSearch}()
+GreedyMethodOptimizer() = TreeOptimizer{:GreedyMethod}()
+KaHyParBipartiteOptimizer() = TreeOptimizer{:KaHyParBipartite}()
+TreeSAOptimizer() = TreeOptimizer{:TreeSA}()
+SABipartiteOptimizer() = TreeOptimizer{:SABipartite}()
+ExactTreewidthOptimizer() = TreeOptimizer{:ExactTreewidth}()
 
 function optimaltree(network, optdata::Dict;
                      optimizer::TreeOptimizer{T}=TreeOptimizer{:ExhaustiveSearch}(),
diff --git a/test/omeinsumcontractionordres.jl b/test/omeinsum.jl
similarity index 77%
rename from test/omeinsumcontractionordres.jl
rename to test/omeinsum.jl
index f8280e38..68b6682c 100644
--- a/test/omeinsumcontractionordres.jl
+++ b/test/omeinsum.jl
@@ -15,8 +15,6 @@ end
 
 using OMEinsumContractionOrders
 
-# open the debug mode to check the optimization algorithms used
-ENV["JULIA_DEBUG"] = "TensorOperationsOMEinsumContractionOrdersExt"
 
 @testset "OMEinsumContractionOrders optimization algorithms" begin
     A = randn(5, 5, 5, 5)
@@ -140,12 +138,28 @@ end
     end
 
     D2 = ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]])
-    D3 = ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]], :ExhaustiveSearch)
-    D4 = ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]], :GreedyMethod)
-    D5 = ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]], :KaHyParBipartite)
-    D6 = ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]], :TreeSA)
-    D7 = ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]], :SABipartite)
-    D8 = ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]], :ExactTreewidth)
+    D3 = ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]], optimizer = ExhaustiveSearchOptimizer())
+    D4 = ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]], optimizer = GreedyMethodOptimizer())
+    D5 = ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]], optimizer = KaHyParBipartiteOptimizer())
+    D6 = ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]], optimizer = TreeSAOptimizer())
+    D7 = ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]], optimizer = SABipartiteOptimizer())
+    D8 = ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]], optimizer = ExactTreewidthOptimizer())
 
     @test D1 ≈ D2 ≈ D3 ≈ D4 ≈ D5 ≈ D6 ≈ D7 ≈ D8
+
+    @test_throws ArgumentError begin
+        D9 = ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]], order = [5, 6, 7], optimizer = GreedyMethod())
+    end
+
+    @test_logs (:debug, "Using optimizer ExhaustiveSearch") min_level=Logging.Debug match_mode=:any ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]], optimizer = ExhaustiveSearchOptimizer())
+
+    @test_logs (:debug, "Using optimizer GreedyMethod from OMEinsumContractionOrders") min_level=Logging.Debug match_mode=:any ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]], optimizer = GreedyMethodOptimizer())
+
+    @test_logs (:debug, "Using optimizer KaHyParBipartite from OMEinsumContractionOrders") min_level=Logging.Debug match_mode=:any ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]], optimizer = KaHyParBipartiteOptimizer())
+
+    @test_logs (:debug, "Using optimizer TreeSA from OMEinsumContractionOrders") min_level=Logging.Debug match_mode=:any ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]], optimizer = TreeSAOptimizer())
+
+    @test_logs (:debug, "Using optimizer SABipartite from OMEinsumContractionOrders") min_level=Logging.Debug match_mode=:any ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]], optimizer = SABipartiteOptimizer())
+    
+    @test_logs (:debug, "Using optimizer ExactTreewidth from OMEinsumContractionOrders") min_level=Logging.Debug match_mode=:any ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]], optimizer = ExactTreewidthOptimizer())
 end
diff --git a/test/runtests.jl b/test/runtests.jl
index abef21f0..44860e21 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1,6 +1,7 @@
 using TensorOperations
 using LinearAlgebra
 using Test
+using Logging
 using Random
 Random.seed!(1234567)
 
@@ -41,8 +42,11 @@ end
 
 # note: OMEinsumContractionOrders should not be loaded before this point
 # as there is a test which requires it to be loaded after
-@testset "OMEinsumOptimizer extension" begin
-    include("omeinsumcontractionordres.jl")
+# it only support julia version >= 1.9
+if VERSION >= v"1.9"
+    @testset "OMEinsumOptimizer extension" begin
+        include("omeinsum.jl")
+    end
 end
 
 @testset "Polynomials" begin

From 4adc9cf9e2668d49e7f26230d7abbcdbfb86c9dc Mon Sep 17 00:00:00 2001
From: ArrogantGao <gaoxuanzhao@gmail.com>
Date: Thu, 15 Aug 2024 12:04:58 +0800
Subject: [PATCH 13/15] format src

---
 src/TensorOperations.jl    | 4 +++-
 src/implementation/ncon.jl | 4 ++--
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/src/TensorOperations.jl b/src/TensorOperations.jl
index 583f9853..cf773983 100644
--- a/src/TensorOperations.jl
+++ b/src/TensorOperations.jl
@@ -30,7 +30,9 @@ export IndexTuple, Index2Tuple, linearize
 export checkcontractible, tensorcost
 
 # export optimizer
-export TreeOptimizer, ExhaustiveSearchOptimizer, GreedyMethodOptimizer, KaHyParBipartiteOptimizer, TreeSAOptimizer, SABipartiteOptimizer, ExactTreewidthOptimizer
+export TreeOptimizer, ExhaustiveSearchOptimizer, GreedyMethodOptimizer,
+       KaHyParBipartiteOptimizer, TreeSAOptimizer, SABipartiteOptimizer,
+       ExactTreewidthOptimizer
 
 # Interface and index types
 #---------------------------
diff --git a/src/implementation/ncon.jl b/src/implementation/ncon.jl
index 42d28f22..89425339 100644
--- a/src/implementation/ncon.jl
+++ b/src/implementation/ncon.jl
@@ -48,7 +48,7 @@ function ncon(tensors, network,
     if isnothing(order)
         if isnothing(optimizer)
             # not specifing order and optimizer, tree via ncontree
-            tree=ncontree(network)
+            tree = ncontree(network)
         else
             # order via optimizer
             optdata = Dict{Any,Number}()
@@ -57,7 +57,7 @@ function ncon(tensors, network,
                     optdata[id] = tensorstructure(tensors[i], j, conjlist[i])
                 end
             end
-            tree=optimaltree(network, optdata, optimizer, false)[1]
+            tree = optimaltree(network, optdata, optimizer, false)[1]
         end
     else
         if !isnothing(optimizer)

From a1206202d56538956016a40f4f0c37443e9885c2 Mon Sep 17 00:00:00 2001
From: ArrogantGao <gaoxuanzhao@gmail.com>
Date: Fri, 16 Aug 2024 12:39:39 +0800
Subject: [PATCH 14/15] update project.toml

---
 Project.toml     | 2 +-
 test/runtests.jl | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/Project.toml b/Project.toml
index 3934e7b2..ca3254fc 100644
--- a/Project.toml
+++ b/Project.toml
@@ -39,7 +39,7 @@ DynamicPolynomials = "0.5"
 LRUCache = "1"
 LinearAlgebra = "1.6"
 Logging = "1.6"
-OMEinsumContractionOrders = "0.9"
+OMEinsumContractionOrders = "0.9.2"
 PackageExtensionCompat = "1"
 PtrArrays = "1.2"
 Random = "1"
diff --git a/test/runtests.jl b/test/runtests.jl
index 44860e21..dc48e2eb 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -42,8 +42,8 @@ end
 
 # note: OMEinsumContractionOrders should not be loaded before this point
 # as there is a test which requires it to be loaded after
-# it only support julia version >= 1.9
-if VERSION >= v"1.9"
+# the tests only work when extension is supported (julia version >= 1.9)
+if isdefined(Base, :get_extension)
     @testset "OMEinsumOptimizer extension" begin
         include("omeinsum.jl")
     end

From d39c6eac39391d22e356692c310799dd99706ce8 Mon Sep 17 00:00:00 2001
From: ArrogantGao <gaoxuanzhao@gmail.com>
Date: Fri, 16 Aug 2024 14:19:02 +0800
Subject: [PATCH 15/15] format code

---
 test/omeinsum.jl | 54 ++++++++++++++++++++++++++++++++++--------------
 1 file changed, 39 insertions(+), 15 deletions(-)

diff --git a/test/omeinsum.jl b/test/omeinsum.jl
index 68b6682c..fc8491c9 100644
--- a/test/omeinsum.jl
+++ b/test/omeinsum.jl
@@ -15,7 +15,6 @@ end
 
 using OMEinsumContractionOrders
 
-
 @testset "OMEinsumContractionOrders optimization algorithms" begin
     A = randn(5, 5, 5, 5)
     B = randn(5, 5, 5)
@@ -138,28 +137,53 @@ end
     end
 
     D2 = ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]])
-    D3 = ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]], optimizer = ExhaustiveSearchOptimizer())
-    D4 = ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]], optimizer = GreedyMethodOptimizer())
-    D5 = ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]], optimizer = KaHyParBipartiteOptimizer())
-    D6 = ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]], optimizer = TreeSAOptimizer())
-    D7 = ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]], optimizer = SABipartiteOptimizer())
-    D8 = ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]], optimizer = ExactTreewidthOptimizer())
+    D3 = ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]];
+              optimizer=ExhaustiveSearchOptimizer())
+    D4 = ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]];
+              optimizer=GreedyMethodOptimizer())
+    D5 = ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]];
+              optimizer=KaHyParBipartiteOptimizer())
+    D6 = ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]];
+              optimizer=TreeSAOptimizer())
+    D7 = ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]];
+              optimizer=SABipartiteOptimizer())
+    D8 = ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]];
+              optimizer=ExactTreewidthOptimizer())
 
     @test D1 ≈ D2 ≈ D3 ≈ D4 ≈ D5 ≈ D6 ≈ D7 ≈ D8
 
     @test_throws ArgumentError begin
-        D9 = ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]], order = [5, 6, 7], optimizer = GreedyMethod())
+        D9 = ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]]; order=[5, 6, 7],
+                  optimizer=GreedyMethod())
+    end
+
+    @test_logs (:debug, "Using optimizer ExhaustiveSearch") min_level = Logging.Debug match_mode = :any begin
+        ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]];
+             optimizer=ExhaustiveSearchOptimizer())
     end
 
-    @test_logs (:debug, "Using optimizer ExhaustiveSearch") min_level=Logging.Debug match_mode=:any ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]], optimizer = ExhaustiveSearchOptimizer())
+    @test_logs (:debug, "Using optimizer GreedyMethod from OMEinsumContractionOrders") min_level = Logging.Debug match_mode = :any begin
+        ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]];
+             optimizer=GreedyMethodOptimizer())
+    end
 
-    @test_logs (:debug, "Using optimizer GreedyMethod from OMEinsumContractionOrders") min_level=Logging.Debug match_mode=:any ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]], optimizer = GreedyMethodOptimizer())
+    @test_logs (:debug, "Using optimizer KaHyParBipartite from OMEinsumContractionOrders") min_level = Logging.Debug match_mode = :any begin
+        ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]];
+             optimizer=KaHyParBipartiteOptimizer())
+    end
 
-    @test_logs (:debug, "Using optimizer KaHyParBipartite from OMEinsumContractionOrders") min_level=Logging.Debug match_mode=:any ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]], optimizer = KaHyParBipartiteOptimizer())
+    @test_logs (:debug, "Using optimizer TreeSA from OMEinsumContractionOrders") min_level = Logging.Debug match_mode = :any begin
+        ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]];
+             optimizer=TreeSAOptimizer())
+    end
 
-    @test_logs (:debug, "Using optimizer TreeSA from OMEinsumContractionOrders") min_level=Logging.Debug match_mode=:any ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]], optimizer = TreeSAOptimizer())
+    @test_logs (:debug, "Using optimizer SABipartite from OMEinsumContractionOrders") min_level = Logging.Debug match_mode = :any begin
+        ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]];
+             optimizer=SABipartiteOptimizer())
+    end
 
-    @test_logs (:debug, "Using optimizer SABipartite from OMEinsumContractionOrders") min_level=Logging.Debug match_mode=:any ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]], optimizer = SABipartiteOptimizer())
-    
-    @test_logs (:debug, "Using optimizer ExactTreewidth from OMEinsumContractionOrders") min_level=Logging.Debug match_mode=:any ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]], optimizer = ExactTreewidthOptimizer())
+    @test_logs (:debug, "Using optimizer ExactTreewidth from OMEinsumContractionOrders") min_level = Logging.Debug match_mode = :any begin
+        ncon([A, B, C], [[-1, 5, -3, 6], [7, -4, 5], [7, 6, -2]];
+             optimizer=ExactTreewidthOptimizer())
+    end
 end