TensorOperations.jl
TensorOperations.jl copied to clipboard
Adding OMEinsumContractionOrders.jl as a backend of TensorOperations.jl for finding the optimal contraction order
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
andSABipartite
: 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 solverTreeWidthSolver.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