Reference

AbstractEinsum

Abstract type for einsum notations.

Required Interfaces

• getixsv: a vector of vectors, each vector represents the labels associated with a input tensor.
• getiyv: a vector of labels associated with the output tensor.
• uniquelabels: a vector of labels that are unique in the einsum notation.

Derived interfaces

• labeltype: the data type to represent the labels in the einsum notation.

CodeOptimizer

Abstract type for code optimizers.

CodeSlicer

Abstract type for code slicers.

EinCode{LT} <: AbstractEinsum
EinCode(ixs::Vector{Vector{LT}}, iy::Vector{LT})

Einsum code with input indices ixs and output index iy.

Examples

The einsum notation for matrix multiplication is:

julia> code = OMEinsumContractionOrders.EinCode([[1,2], [2, 3]], [1, 3])
1∘2, 2∘3 -> 1∘3

julia> OMEinsumContractionOrders.getixsv(code)
2-element Vector{Vector{Int64}}:
 [1, 2]
 [2, 3]

julia> OMEinsumContractionOrders.getiyv(code)
2-element Vector{Int64}:
 1
 3

NestedEinsum{LT} <: AbstractEinsum
NestedEinsum(args::Vector{NestedEinsum}, eins::EinCode)

The einsum notation with a contraction order specified as a tree data structure. It is automatically generated by the contraction code optimizer with the optimize_code function.

Fields

• args: the children of the current node
• tensorindex: the index of the input tensor, required only for leaf nodes. For non-leaf nodes, it is -1.
• eins: the einsum notation for the operation at the current node.

ScoreFunction

A function to compute the score of a contraction code:

score = tc_weight * 2^tc + rw_weight * 2^rw + sc_weight * max(0, 2^sc - 2^sc_target)

Fields

• tc_weight: the weight of the time complexity, default is 1.0.
• sc_weight: the weight of the space complexity (the size of the largest tensor), default is 1.0.
• rw_weight: the weight of the read-write complexity, default is 0.0.
• sc_target: the target space complexity, below which the sc_weight will be set to 0 automatically, default is 0.0.

SlicedEinsum{LT,ET<:Union{EinCode{LT},NestedEinsum{LT}}} <: AbstractEinsum
SlicedEinsum(slicing::Vector{LT}, eins::ET)

The einsum notation with sliced indices. The sliced indices are the indices enumerated manually at the top level. By slicing the indices, the space complexity of the einsum notation can be reduced.

Fields

• slicing: the sliced indices.
• eins: the einsum notation of the current node, which is a NestedEinsum object.

getixsv(code::AbstractEinsum) -> Vector{Vector{LT}}

Returns the input indices of the einsum notation. Each vector represents the labels associated with a input tensor.

getiyv(code::AbstractEinsum) -> Vector{LT}

Returns the output index of the einsum notation.

labeltype(code::AbstractEinsum) -> Type

Returns the data type to represent the labels in the einsum notation.

uniquelabels(code::AbstractEinsum) -> Vector{LT}

Returns the unique labels in the einsum notation. The labels are the indices of the tensors.

contraction_complexity(eincode, size_dict) -> ContractionComplexity

Returns the time, space and read-write complexity of the einsum contraction. The returned ContractionComplexity object contains 3 fields:

• tc: time complexity defined as log2(number of element-wise multiplications).
• sc: space complexity defined as log2(size of the maximum intermediate tensor).
• rwc: read-write complexity defined as log2(the number of read-write operations).

flop(eincode, size_dict) -> Int

Returns the number of iterations, which is different with the true floating point operations (FLOP) by a factor of 2.

label_elimination_order(code) -> Vector

Returns a vector of labels sorted by the order they are eliminated in the contraction tree. The contraction tree is specified by code, which e.g. can be a NestedEinsum instance.

peak_memory(code, size_dict::Dict) -> Int

Estimate peak memory in number of elements.

uniformsize(code::AbstractEinsum, size::Int) -> Dict

Returns a dictionary that maps each label to the given size.

optimize_code(eincode, size_dict, optimizer = GreedyMethod(); slicer=nothing, simplifier=nothing, permute=true) -> optimized_eincode

Optimize the einsum contraction code and reduce the time/space complexity of tensor network contraction. Returns a NestedEinsum instance. Input arguments are

Arguments

• eincode is an einsum contraction code instance, one of DynamicEinCode, StaticEinCode or NestedEinsum.
• size is a dictionary of "edge label=>edge size" that contains the size information, one can use uniformsize(eincode, 2) to create a uniform size.
• optimizer is a CodeOptimizer instance, should be one of GreedyMethod, Treewidth, KaHyParBipartite, SABipartite or TreeSA. Check their docstrings for details.

Keyword Arguments

• slicer is for slicing the contraction code to reduce the space complexity, default is nothing. Currently only TreeSASlicer is supported.
• simplifier is one of MergeVectors or MergeGreedy. Default is nothing.
• permute is a boolean flag to indicate whether to optimize the permutation of the contraction order.

Examples

julia> using OMEinsum

julia> code = ein"ij, jk, kl, il->"
ij, jk, kl, il -> 

julia> optimize_code(code, uniformsize(code, 2), TreeSA());

slice_code(code, size_dict, slicer) -> sliced_code

Slice the einsum contraction code to reduce the space complexity, returns a SlicedEinsum instance.

Arguments

• code is a NestedEinsum instance.
• size_dict is a dictionary of "edge label=>edge size" that contains the size information, one can use uniformsize(eincode, 2) to create a uniform size.
• slicer is a CodeSlicer instance, currently only TreeSASlicer is supported.

GreedyMethod{MT}
GreedyMethod(; α = 0.0, temperature = 0.0)

It may not be optimal, but it is fast.

Fields

• α is the parameter for the loss function, for pairwise interaction, L = size(out) - α * (size(in1) + size(in2))
• temperature is the parameter for sampling, if it is zero, the minimum loss is selected; for non-zero, the loss is selected by the Boltzmann distribution, given by p ~ exp(-loss/temperature).

optimize_greedy(eincode, size_dict; α, temperature)

Greedy optimizing the contraction order and return a NestedEinsum object. Check the docstring of tree_greedy for detailed explaination of other input arguments.

tree_greedy(incidence_list, log2_sizes; α = 0.0, temperature = 0.0)

Compute greedy order, and the time and space complexities, the rows of the incidence_list are vertices and columns are edges. log2_sizes are defined on edges. α is the parameter for the loss function, for pairwise interaction, L = size(out) - α * (size(in1) + size(in2)) temperature is the parameter for sampling, if it is zero, the minimum loss is selected; for non-zero, the loss is selected by the Boltzmann distribution, given by p ~ exp(-loss/temperature).

TreeSA{IT} <: CodeOptimizer
TreeSA(; βs=collect(0.01:0.05:15), ntrials=10, niters=50, initializer=:greedy, score=ScoreFunction())

Optimize the einsum contraction pattern using the simulated annealing on tensor expression tree.

Fields

• ntrials, βs and niters are annealing parameters, doing ntrials indepedent annealings, each has inverse tempteratures specified by βs, in each temperature, do niters updates of the tree.
• initializer specifies how to determine the initial configuration, it can be :greedy, :random or :specified. If the initializer is :specified, the input code should be a NestedEinsum object.
• score specifies the score function to evaluate the quality of the contraction tree, it is a function of time complexity, space complexity and read-write complexity.

References

• Recursive Multi-Tensor Contraction for XEB Verification of Quantum Circuits

Breaking changes:

• nslices is removed, since the slicing part is now separated from the optimization part, see slice_code function and TreeSASlicer.
• greedy_method is removed. If you want to have detailed control of the initializer, please pre-optimize the code with another method and then use :specified to initialize the tree.

optimize_tree(code, size_dict; βs, ntrials, niters, initializer, score)

Optimize the einsum contraction pattern specified by code, and edge sizes specified by size_dict. Check the docstring of TreeSA for detailed explaination of other input arguments.

const ExactTreewidth = Treewidth{SafeRules{BT, MMW{3}(), MF}}
ExactTreewidth() = Treewidth()

ExactTreewidth is a specialization of Treewidth for the SafeRules preprocessing algorithm with the BT elimination algorithm. The BT algorithm is an exact solver for the treewidth problem that implemented in TreeWidthSolver.jl.

struct Treewidth{EL <: EliminationAlgorithm, GM} <: CodeOptimizer
Treewidth(; alg::EL = SafeRules(BT(), MMW{3}(), MF()))

Tree width based solver. The solvers are implemented in CliqueTrees.jl and TreeWidthSolver.jl. They include:

Detailed descriptions is available in the CliqueTrees.jl.

Fields

• alg::EL: The algorithm to use for the treewidth calculation. Available elimination algorithms are listed above.

Example

julia> optimizer = Treewidth();

julia> eincode = OMEinsumContractionOrders.EinCode([['a', 'b'], ['a', 'c', 'd'], ['b', 'c', 'e', 'f'], ['e'], ['d', 'f']], ['a'])
ab, acd, bcef, e, df -> a

julia> size_dict = Dict([c=>(1<<i) for (i,c) in enumerate(['a', 'b', 'c', 'd', 'e', 'f'])]...)
Dict{Char, Int64} with 6 entries:
  'f' => 64
  'a' => 2
  'c' => 8
  'd' => 16
  'e' => 32
  'b' => 4

julia> optcode = optimize_code(eincode, size_dict, optimizer)
ba, ab -> a
├─ bcf, fac -> ba
│  ├─ e, bcef -> bcf
│  │  ├─ e
│  │  └─ bcef
│  └─ df, acd -> fac
│     ├─ df
│     └─ acd
└─ ab

optimize_treewidth(optimizer, eincode, size_dict)

Optimizing the contraction order via solve the exact tree width of the line graph corresponding to the eincode and return a NestedEinsum object. Check the docstring of treewidth_method for detailed explaination of other input arguments.

HyperND(;
    dis = KaHyParND(),
    algs = (MF(), AMF(), MMD()),
    level = 6,
    width = 120,
    imbalances = 130:130,
    score = ScoreFunction(),
)

Nested-dissection based optimizer. Recursively partitions a tensor network, then calls a greedy algorithm on the leaves. The optimizer is run a number of times: once for each greedy algorithm in algs and each imbalance value in imbalances. The recursion depth is controlled by the parameters level and width.

The line graph is partitioned using the algorithm dis. OMEinsumContractionOrders currently supports two partitioning algorithms, both of which require importing an external library.

The optimizer is implemented using the tree decomposition library CliqueTrees.jl.

Arguments

• dis: graph partitioning algorithm
• algs: tuple of elimination algorithms.
• level: maximum level
• width: minimum width
• imbalances: imbalance parameters
• score: a function to evaluate the quality of the contraction tree. Default is ScoreFunction().

BipartiteResult{RT}
BipartiteResult(part1, part2, sc, valid)

Result of the bipartite optimization. part1 and part2 are the two parts of the bipartition, sc is the space complexity of the bipartition, valid is a boolean indicating whether the bipartition is valid.

KaHyParBipartite{RT,IT,GM}
KaHyParBipartite(; sc_target, imbalances=collect(0.0:0.005:0.8),
    max_group_size=40, greedy_config=GreedyMethod())

Optimize the einsum code contraction order using the KaHyPar + Greedy approach. This program first recursively cuts the tensors into several groups using KaHyPar, with maximum group size specifed by max_group_size and maximum space complexity specified by sc_target, Then finds the contraction order inside each group with the greedy search algorithm. Other arguments are

Fields

• sc_target is the target space complexity, defined as log2(number of elements in the largest tensor),
• imbalances is a KaHyPar parameter that controls the group sizes in hierarchical bipartition,
• max_group_size is the maximum size that allowed to used greedy search,
• sub_optimizer is the sub-optimizer used to find the contraction order when the group size is small enough.

References

• Hyper-optimized tensor network contraction
• Simulating the Sycamore quantum supremacy circuits

SABipartite{RT,BT}
SABipartite(; sc_target=25, ntrials=50, βs=0.1:0.2:15.0, niters=1000
    max_group_size=40, greedy_config=GreedyMethod(), initializer=:random)

Optimize the einsum code contraction order using the Simulated Annealing bipartition + Greedy approach. This program first recursively cuts the tensors into several groups using simulated annealing, with maximum group size specifed by max_group_size and maximum space complexity specified by sc_target, Then finds the contraction order inside each group with the greedy search algorithm. Other arguments are

Fields

• sc_target is the target space complexity, defined as log2(number of elements in the largest tensor),
• ntrials is the number of repetition (with different random seeds),
• βs is a list of inverse temperature 1/T,
• niters is the number of iteration in each temperature,
• max_group_size is the maximum size that allowed to used greedy search,
• sub_optimizer is the optimizer for the bipartited sub graphs, one can choose GreedyMethod() or TreeSA(),
• initializer is the partition configuration initializer, one can choose :random or :greedy (slow but better).

References

• Hyper-optimized tensor network contraction

optimize_sa(code, size_dict; sc_target, max_group_size=40, βs=0.1:0.2:15.0, niters=1000, ntrials=50,
       sub_optimizer = GreedyMethod(), initializer=:random)

Optimize the einsum code contraction order using the Simulated Annealing bipartition + Greedy approach. size_dict is a dictionary that specifies leg dimensions. Check the docstring of SABipartite for detailed explaination of other input arguments.

References

• Hyper-optimized tensor network contraction

TreeSASlicer{IT, LT} <: CodeSlicer

A structure for configuring the Tree Simulated Annealing (TreeSA) slicing algorithm. The goal of slicing is to reach the target space complexity specified by score.sc_target.

Fields

• ntrials, βs and niters are annealing parameters, doing ntrials indepedent annealings, each has inverse tempteratures specified by βs, in each temperature, do niters updates of the tree.
• fixed_slices::Vector{LT}: A vector of fixed slices that should not be altered. Default is an empty vector.
• optimization_ratio::Float64: A constant used for determining the number of iterations for slicing. Default is 2.0. i.e. if the current space complexity is 30, and the target space complexity is 20, then the number of iterations for slicing is (30 - 20) x optimization_ratio.
• score::ScoreFunction: A function to evaluate the quality of the contraction tree. Default is ScoreFunction(sc_target=30.0).

References

• Recursive Multi-Tensor Contraction for XEB Verification of Quantum Circuits

CodeSimplifier

Abstract type for code simplifiers.

MergeGreedy <: CodeSimplifier
MergeGreedy(; threshhold=-1e-12)

Contraction code simplifier (in order to reduce the time of calling optimizers) that merges tensors greedily if the space complexity of merged tensors is reduced (difference smaller than the threshhold).

MergeVectors <: CodeSimplifier
MergeVectors()

Contraction code simplifier (in order to reduce the time of calling optimizers) that merges vectors to closest tensors.

NetworkSimplifier{LT}

A network simplifier that contains a list of operations that can be applied to a tensor network to reduce the number of tensors. It is generated from a proprocessor, such as MergeVectors or MergeGreedy.

Fields

• operations: a list of NestedEinsum objects.

embed_simplifier(code::NestedEinsum, simplifier::NetworkSimplifier)

Embed the simplifier into the contraction code. A typical workflow is: (i) generate a simplifier with simplify_code, (ii) then optimize the simplified code with optimize_code and (iii) post-process the optimized code with embed_simplifier to produce correct contraction order for the original code. This is automatically done in optimize_code given the simplifier argument is not nothing.

Arguments

• code: the contraction code to embed the simplifier into.
• simplifier: the simplifier to embed, which is a NetworkSimplifier object.

Returns

• A new NestedEinsum object.

simplify_code(code::Union{EinCode, NestedEinsum}, size_dict, method::CodeSimplifier)

Simplify the contraction code by preprocessing the code with a simplifier.

Arguments

• code: the contraction code to simplify.
• size_dict: the size dictionary of the contraction code.
• method: the simplifier to use, which can be MergeVectors or MergeGreedy.

Returns

• A tuple of (NetworkSimplifier, newcode), where newcode is a new EinCode object.

readjson(filename::AbstractString)

Read the contraction order from a JSON file.

Arguments

• filename: the name of the file to read from.

writejson(filename::AbstractString, ne::Union{NestedEinsum, SlicedEinsum})

Write the contraction order to a JSON file.

Arguments

• filename: the name of the file to write to.
• ne: the contraction order to write. It can be a NestedEinsum or a SlicedEinsum object.

viz_contraction(code::Union{NestedEinsum, SlicedEinsum}; locs=StressLayout(), framerate=10, filename=tempname() * ".mp4", show_progress=true)

Visualize the contraction process of a tensor network.

Arguments

• code: The tensor network to visualize.

Keyword Arguments

• locs: The coordinates or layout algorithm to use for positioning the nodes in the graph. Default is StressLayout().
• framerate: The frame rate of the animation. Default is 10.
• filename: The name of the output file, with .gif or .mp4 extension. Default is a temporary file with .mp4 extension.
• show_progress: Whether to show progress information. Default is true.

Returns

• the path of the generated file.

viz_eins(code::AbstractEinsum; locs=StressLayout(), filename = nothing, kwargs...)

Visualizes an AbstractEinsum object by creating a tensor network graph and rendering it using GraphViz.

Arguments

• code::AbstractEinsum: The AbstractEinsum object to visualize.

Keyword Arguments

• locs=StressLayout(): The coordinates or layout algorithm to use for positioning the nodes in the graph.
• filename = nothing: The name of the file to save the visualization to. If nothing, the visualization will be displayed on the screen instead of saving to a file.
• config = GraphDisplayConfig(): The configuration for displaying the graph. Please refer to the documentation of GraphDisplayConfig for more information.
• kwargs...: Additional keyword arguments to be passed to the GraphViz constructor.