TensorOperations.jl icon indicating copy to clipboard operation
TensorOperations.jl copied to clipboard

Adding OMEinsumContractionOrders.jl as a backend of TensorOperations.jl for finding the optimal contraction order

Open ArrogantGao opened this issue 6 months ago • 9 comments

An extension OMEinsumContractionOrdersExt has been added, which provides OMEinsumContractionOrders.jl as a backend of optimizing the contraction order.

We modified the @tensor marco and added a keyword opt_algorithm for selecting a specific optimizer. The default optimizer is NCon, which is the original one used in TensorOperations.jl. After the extension is loaded, the following optimizers are provided:

  • GreedyMethod: greedily select the pair of tensors with the smallest cost to contract at each step
  • TreeSA: tree simulated annealing, based on local search and simulating annealing
  • KaHyParBipartite and SABipartite: regarding the tensor network as hypergraph and then partition the hypergraph into two parts, and then recursively partition each part
  • ExactTreewidth: method based on exact tree width solver TreeWidthSolver.jl, using the tree decomposition with minimal tree width to construct contraction order with optimal complexity The former four optimizers are approximate ones which are suitable for large networks, while the latest one is an exact solver, which is only suitable for small networks.

Here is an example of use

using TensorOperations, OMEinsumContractionOrders
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 = (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

ArrogantGao avatar Aug 07 '24 15:08 ArrogantGao