Contraction order optimization

Constructing a code

The @ein_str string literal does not optimize the contraction order for more than two input tensors. We first use a graph to construct a DynamicEinCode object for demonstration.

julia> using OMEinsum, OMEinsumContractionOrders, OMEinsumContractionOrders.Graphs
julia> graph = random_regular_graph(20, 3; seed=42){20, 30} undirected simple Int64 graph
julia> code = EinCode([[e.src, e.dst] for e in edges(graph)], Int[])1∘2, 1∘11, 1∘15, 2∘19, 2∘20, 3∘4, 3∘6, 3∘18, 4∘13, 4∘17, 5∘8, 5∘10, 5∘20, 6∘13, 6∘19, 7∘8, 7∘10, 7∘12, 8∘14, 9∘12, 9∘14, 9∘17, 10∘15, 11∘13, 11∘18, 12∘17, 14∘20, 15∘16, 16∘18, 16∘19 ->

The return value is a StaticEinCode object that does not contain a contraction order. The time and space complexity can be obtained by calling the contraction_complexity function.

julia> size_dict = uniformsize(code, 2)  # size of the labels are set to 2Dict{Int64, Int64} with 20 entries:
  5  => 2
  16 => 2
  20 => 2
  7  => 2
  12 => 2
  8  => 2
  17 => 2
  1  => 2
  19 => 2
  4  => 2
  15 => 2
  6  => 2
  2  => 2
  11 => 2
  18 => 2
  13 => 2
  10 => 2
  9  => 2
  14 => 2
  3  => 2
julia> contraction_complexity(code, size_dict)  # time and space complexityTime complexity: 2^20.0
Space complexity: 2^0.0
Read-write complexity: 2^6.918863237274595

The return values are log2 values of the number of iterations, number of elements of the largest tensor and the number of elementwise read-write operations.

Optimizing the contraction order

To optimize the contraction order, we can use the optimize_code function.

julia> optcode = optimize_code(code, size_dict, TreeSA(ntrials=1))7∘10, 7∘10 ->
├─ 7∘15∘5, 15∘5∘10 -> 7∘10
│  ├─ 5∘4∘14∘7, 4∘15∘5∘14 -> 7∘15∘5
│  │  ├─ 5∘14∘7, 7∘14∘4 -> 5∘4∘14∘7
│  │  │  ├─ 5∘14∘8, 7∘8 -> 5∘14∘7
│  │  │  │  ├─ 5∘8, 8∘14 -> 5∘14∘8
│  │  │  │  │  ⋮
│  │  │  │  │
│  │  │  │  └─ 7∘8
│  │  │  └─ 7∘14∘17, 4∘17 -> 7∘14∘4
│  │  │     ├─ 17∘7∘12, 14∘17∘12 -> 7∘14∘17
│  │  │     │  ⋮
│  │  │     │
│  │  │     └─ 4∘17
│  │  └─ 4∘15∘20, 5∘14∘20 -> 4∘15∘5∘14
│  │     ├─ 4∘15∘2, 2∘20 -> 4∘15∘20
│  │     │  ├─ 2∘4∘11∘16, 16∘11∘2∘15 -> 4∘15∘2
│  │     │  │  ⋮
│  │     │  │
│  │     │  └─ 2∘20
│  │     └─ 5∘20, 14∘20 -> 5∘14∘20
│  │        ├─ 5∘20
│  │        └─ 14∘20
│  └─ 10∘15, 5∘10 -> 15∘5∘10
│     ├─ 10∘15
│     └─ 5∘10
└─ 7∘10

The output value is a binary contraction tree with type SlicedEinsum or NestedEinsum. The TreeSA is a local search algorithm that optimizes the contraction order. More algorithms can be found in the OMEinsumContractionOrders documentation.

After optimizing the contraction order, the time and readwrite complexities are significantly reduced.

julia> contraction_complexity(optcode, size_dict)Time complexity: 2^8.820178962415188
Space complexity: 2^4.0
Read-write complexity: 2^9.359749560322332

Slicing the code

In some cases, the memory usage of the contraction is too large. Slicing is a technique to reduce the time and space complexity of the contraction. The slicing is done by using the slice_code function.

julia> slicer = TreeSASlicer(score=ScoreFunction(sc_target=2))TreeSASlicer{StepRangeLen{Float64, Base.TwicePrecision{Float64}, Base.TwicePrecision{Float64}, Int64}, Any}(14.0:0.05:15.0, 10, 10, Any[], 2.0, ScoreFunction(1.0, 1.0, 0.0, 2.0))
julia> scode = slice_code(optcode, size_dict, slicer);
julia> contraction_complexity(scode, size_dict)Time complexity: 2^10.321928094887362
Space complexity: 2^2.0
Read-write complexity: 2^11.262094845370179
julia> scode.slicing3-element Vector{Int64}:
 19
 20
  4

The return value is a SlicedEinsum object. The space complexity is reduced to 2, while the time complexity is increased as a trade-off.

Using optein string literal

For convenience, the optimized contraction can be directly contructed by using the @optein_str string literal.

julia> optein"ij,jk,kl,li->"  # optimized contraction, without knowing the size of the tensorsik, ik ->
├─ li, kl -> ik
│  ├─ li
│  └─ kl
└─ ij, jk -> ik
   ├─ ij
   └─ jk

@optein_str optimizes the contraction order with the assumption that each index has the same size 2, hence the resulting contraction order might not be optimal.

Manual optimization

One can also manually specify the contraction order by using the @ein_str string literal.

julia> ein"((ij,jk),kl),li->ik"  # manually optimized contractionikl, li -> ik
├─ ik, kl -> ikl
│  ├─ ij, jk -> ik
│  │  ├─ ij
│  │  └─ jk
│  └─ kl
└─ li

Flatten the code

Given an optimized code, one can flatten it to get a code without contraction order with type EinCode.

julia> OMEinsum.flatten(optcode)1∘2, 1∘11, 1∘15, 2∘19, 2∘20, 3∘4, 3∘6, 3∘18, 4∘13, 4∘17, 5∘8, 5∘10, 5∘20, 6∘13, 6∘19, 7∘8, 7∘10, 7∘12, 8∘14, 9∘12, 9∘14, 9∘17, 10∘15, 11∘13, 11∘18, 12∘17, 14∘20, 15∘16, 16∘18, 16∘19 ->