OMEinsum.DynamicEinCode — TypeDynamicEinCode{LT}
DynamicEinCode(ixs, iy)Wrapper to eincode-specification that creates a callable object to evaluate the eincode ixs -> iy where ixs are the index-labels of the input-tensors and iy are the index-labels of the output.
example
julia> a, b = rand(2,2), rand(2,2);
julia> OMEinsum.DynamicEinCode((('i','j'),('j','k')),('i','k'))(a, b) ≈ a * b
trueOMEinsum.DynamicNestedEinsum — TypeDynamicNestedEinsum{LT} <: NestedEinsum{LT}
DynamicNestedEinsum(args, eins)
DynamicNestedEinsum{LT}(tensorindex::Int)Einsum with contraction order, where the type parameter LT is the label type. It has two constructors. One takes a tensorindex as input, which represents the leaf node in a contraction tree. The other takes an iterable of type DynamicNestedEinsum, args, as the siblings, and eins to specify the contraction operation.
OMEinsum.EinArray — TypeEinArray{T, N, TT, LX, LY, ICT, OCT} <: AbstractArray{T, N}A struct to hold the intermediate result of an einsum where all index-labels of both input and output are expanded to a rank-N-array whose values are lazily calculated. Indices are arranged as inner indices (or reduced dimensions) first and then outer indices.
Type parameters are
* `T`: element type,
* `N`: array dimension,
* `TT`: type of "tuple of input arrays",
* `LX`: type of "tuple of input indexers",
* `LX`: type of output indexer,
* `ICT`: typeof inner CartesianIndices,
* `OCT`: typeof outer CartesianIndices,OMEinsum.EinCode — TypeEinCode <: AbstractEinsum
EinCode(ixs, iy)Abstract type for sum-product contraction code. The constructor returns a DynamicEinCode instance.
OMEinsum.EinIndexer — TypeEinIndexer{locs,N}A structure for indexing EinArrays. locs is the index positions (among all indices). In the constructor, size is the size of target tensor,
OMEinsum.EinIndexer — MethodEinIndexer{locs}(size::Tuple)Constructor for EinIndexer for an object of size size where locs are the locations of relevant indices in a larger tuple.
OMEinsum.IndexGroup — TypeIndexGroupLeaf in a contractiontree, contains the indices and the number of the tensor it describes, e.g. in "ij,jk -> ik", indices "ik" belong to tensor 1, so would be described by IndexGroup(['i','k'], 1).
OMEinsum.NestedEinsum — TypeNestedEinsum{LT} <: AbstractEinsumThe abstract type for contraction trees. It has two subtypes, DynamicNestedEinsum and StaticNestedEinsum.
OMEinsum.NestedEinsumConstructor — TypeNestedEinsumConstructordescribes a (potentially) nested einsum. Important fields:
args, vector of all inputs, eitherIndexGroupobjects corresponding to tensors orNestedEinsumConstructoriy, indices of output
OMEinsum.SlicedEinsum — TypeSlicedEinsum{LT, Ein} <: AbstractEinsumA tensor network with slicing. LT is the label type and Ein is the tensor network.
Fields
slicing::Vector{LT}: A vector of labels to slice.eins::Ein: The tensor network.
OMEinsum.StaticEinCode — TypeStaticEinCode{LT, ixs, iy}The static version of DynamicEinCode that matches the contraction rule at compile time. It is the default return type of @ein_str macro. LT is the label type.
OMEinsum.StaticNestedEinsum — TypeStaticNestedEinsum{LT,args,eins} <: NestedEinsum{LT}
StaticNestedEinsum(args, eins)
StaticNestedEinsum{LT}(tensorindex::Int)Einsum with contraction order, where the type parameter LT is the label type, args is a tuple of StaticNestedEinsum, eins is a StaticEinCode and leaf node is defined by setting eins to an integer. It has two constructors. One takes a tensorindex as input, which represents the leaf node in a contraction tree. The other takes an iterable of type DynamicNestedEinsum, args, as the siblings, and eins to specify the contraction operation.
Base.getindex — Methodgetindex(A::EinArray, inds...)return the lazily calculated entry of A at index inds.
OMEinsum.allow_loops — Methodallow_loops(flag::Bool)Setting this to false will cause OMEinsum to log an error if it falls back to loop_einsum evaluation, instead of calling specialised kernels. The default is true.
OMEinsum.allunique — Methodallunique(ix::Tuple)return true if all elements of ix appear only once in ix.
example
julia> using OMEinsum: allunique
julia> allunique((1,2,3,4))
true
julia> allunique((1,2,3,1))
falseOMEinsum.analyze_binary — MethodGet the expected labels.
OMEinsum.asarray — Methodasarray(x[, parent::AbstractArray]) -> AbstactArrayReturn a 0-dimensional array with item x, otherwise, do nothing. If a parent is supplied, it will try to match the parent array type.
OMEinsum.back_propagate — Methodback_propagate(f, code, cache, ȳ, size_dict)Back propagate the message ȳ through the cached tree cache and return a tree storing the intermediate messages. The message can be gradients et al.
Arguments
f: The back-propagation rule. The signature isf(eins, xs, y, size_dict, dy) -> dxs, whereeins: The contraction code at the current node.xs: The input tensors at the current node.y: The output tensor at the current node.size_dict: The size dictionary, which maps the label to the size of the corresponding dimension.dy: The message on the output tensor (y) to back-propagate through the current node.dxs: The message on the input tensors (xs) as the result of back-propagation.
code: The contraction code, which can be aNestedEinsumor aSlicedEinsum.cache: The cached intermediate results, which can be generated bycached_einsum.ȳ: The message to back-propagate.size_dict: The size dictionary, which maps the label to the size of the corresponding dimension.
Returns
CacheTree: The tree storing the intermediate messages.
OMEinsum.cached_einsum — Methodcached_einsum(code, xs, size_dict)Compute the einsum contraction and cache the intermediate contraction results.
Arguments
code: The contraction code, which can be aNestedEinsumor aSlicedEinsum.xs: The input tensors.size_dict: The size dictionary, which maps the label to the size of the corresponding dimension.
Returns
CacheTree: The cached tree storing the intermediate results.
OMEinsum.cost_and_gradient — Functioncost_and_gradient(code, xs, ȳ)Compute the cost and the gradients w.r.t the input tensors xs.
Arguments
code: The contraction code, which can be aNestedEinsumor aSlicedEinsum.xs: The input tensors.ȳ: The message to back-propagate. Default is1.
Returns
cost: The cost of the contraction.grads: The gradients w.r.t the input tensors.
OMEinsum.einarray — Methodeinarray(::Val{ixs}, Val{iy}, xs, size_dict) -> EinArrayConstructor of EinArray from an EinCode, a tuple of tensors xs and a size_dict that assigns each index-label a size. The returned EinArray holds an intermediate result of the einsum specified by the EinCode with indices corresponding to all unique labels in the einsum. Reduction over the (lazily calculated) dimensions that correspond to labels not present in the output lead to the result of the einsum.
example
julia> using OMEinsum: get_size_dict
julia> a, b = rand(2,2), rand(2,2);
julia> sd = get_size_dict((('i','j'),('j','k')), (a, b));
julia> ea = OMEinsum.einarray(Val((('i','j'),('j','k'))),Val(('i','k')), (a,b), sd);
julia> dropdims(sum(ea, dims=1), dims=1) ≈ a * b
trueOMEinsum.einsum — Functioneinsum(code::EinCode, xs, size_dict)Return the tensor that results from contracting the tensors xs according to the contraction code code.
Arguments
code: The einsum notation, which can be an instance ofEinCode,NestedEinsum, orSlicedEinsum.xs- the input tensorssize_dict- a dictionary that maps index-labels to their sizes
Examples
julia> a, b = rand(2,2), rand(2,2);
julia> einsum(EinCode((('i','j'),('j','k')),('i','k')), (a, b)) ≈ a * b
true
julia> einsum(EinCode((('i','j'),('j','k')),('k','i')), (a, b)) ≈ permutedims(a * b, (2,1))
trueOMEinsum.einsum! — Functioneinsum!(code::EinCode, xs, y, sx, sy, size_dict)Inplace version of einsum. The result is stored in y.
Arguments
code: The einsum notation, which can be an instance ofEinCode,NestedEinsum, orSlicedEinsum.xs: The input tensors.y: The output tensor.sx: Scalexbysx.sy: Scaleybysy.size_dict: A dictionary that maps index-labels to their sizes.
OMEinsum.einsum_grad — Methodeinsum_grad(ixs, xs, iy, size_dict, cdy, i)return the gradient of the result of evaluating the EinCode w.r.t the ith tensor in xs. cdy is the result of applying the EinCode to the xs.
example
julia> using OMEinsum: einsum_grad, get_size_dict
julia> a, b = rand(2,2), rand(2,2);
julia> c = einsum(EinCode((('i','j'),('j','k')), ('i','k')), (a,b));
julia> sd = get_size_dict((('i','j'),('j','k')), (a,b));
julia> einsum_grad((('i','j'),('j','k')), (a,b), ('i','k'), sd, c, 1) ≈ c * transpose(b)
trueOMEinsum.filliys! — Methodfilliys!(neinsum::NestedEinsumConstructor)goes through all NestedEinsumConstructor objects in the tree and saves the correct iy in them.
OMEinsum.get_size_dict! — Methodget_size_dict!(ixs, xs, size_info)return a dictionary that is used to get the size of an index-label in the einsum-specification with input-indices ixs and tensors xs after consistency within ixs and between ixs and xs has been verified.
OMEinsum.getixsv — Methodgetixsv(code)Get labels of input tensors for EinCode, NestedEinsum and some other einsum like objects. Returns a vector of vectors.
julia> getixsv(ein"(ij,jk),k->i")
3-element Vector{Vector{Char}}:
['i', 'j']
['j', 'k']
['k']OMEinsum.getiyv — Methodgetiy(code)Get labels of the output tensor for EinCode, NestedEinsum and some other einsum like objects. Returns a vector.
julia> getiyv(ein"(ij,jk),k->i")
1-element Vector{Char}:
'i': ASCII/Unicode U+0069 (category Ll: Letter, lowercase)OMEinsum.indices_and_locs — Methodindices_and_locs(ixs,iy)given the index-labels of input and output of an einsum, return (in the same order):
- a tuple of the distinct index-labels of the output
iy - a tuple of the distinct index-labels in
ixsof the input not appearing in the outputiy - a tuple of tuples of locations of an index-label in the
ixsin a list of all index-labels - a tuple of locations of index-labels in
iyin a list of all index-labels
where the list of all index-labels is simply the first and the second output catenated and the second output catenated.
OMEinsum.loop_einsum! — Methodloop_einsum!(ixs, iy, xs, y, sx, sy, size_dict)inplace-version of loop_einsum, saving the result in a preallocated tensor of correct size y.
OMEinsum.loop_einsum — Methodloop_einsum(::EinCode, xs, size_dict)evaluates the eincode specified by EinCode and the tensors xs by looping over all possible indices and calculating the contributions ot the result. Scales exponentially in the number of distinct index-labels.
OMEinsum.map_prod — Methodmap_prod(xs, ind, indexers)calculate the value of an EinArray with EinIndexers indexers at location ind.
OMEinsum.match_rule — Methodmatch_rule(ixs, iy)
match_rule(code::EinCode)Returns the rule that matches, otherwise use DefaultRule - the slow loop_einsum backend.
OMEinsum.nopermute — Methodnopermute(ix,iy)check that all values in iy that are also in ix have the same relative order,
example
julia> using OMEinsum: nopermute
julia> nopermute((1,2,3),(1,2))
true
julia> nopermute((1,2,3),(2,1))
falsee.g. nopermute((1,2,3),(1,2)) is true while nopermute((1,2,3),(2,1)) is false
OMEinsum.parse_parens — Methodparse_parens(s::AbstractString, i, narg)parse one level of parens starting at index i where narg counts which tensor the current group of indices, e.g. "ijk", belongs to. Recursively calls itself for each new opening paren that's opened.
OMEinsum.tensorpermute! — Methodtensorpermute(A, perm)permutedims(A, perm) with grouped dimensions.
OMEinsum.@ein! — Macro@ein! A[i,k] := B[i,j] * C[j,k] # A = B * C
@ein! A[i,k] += B[i,j] * C[j,k] # A += B * C
@ein! A[i,k] -= B[i,j] * C[j,k] # A -= B * CMacro interface similar to that of other packages.
Inplace version of @ein.
example
julia> a, b, c, d = rand(2,2), rand(2,2), rand(2,2), zeros(2,2);
julia> cc = copy(c);
julia> @ein! d[i,k] := a[i,j] * b[j,k];
julia> d ≈ a * b
true
julia> d ≈ ein"ij,jk -> ik"(a,b)
true
julia> @ein! c[i,k] += a[i,j] * b[j,k];
julia> c ≈ cc + a * b
true
julia> @ein! c[i,k] -= a[i,j] * b[j,k];
julia> c ≈ cc
trueOMEinsum.@ein — Macro@ein A[i,k] := B[i,j] * C[j,k] # A = B * CMacro interface similar to that of other packages.
You may use numbers in place of letters for dummy indices, as in @tensor, and need not name the output array. Thus A = @ein [1,2] := B[1,ξ] * C[ξ,2] is equivalent to the above. This can also be written A = ein"ij,jk -> ik"(B,C) using the numpy-style string macro.
example
julia> a, b = rand(2,2), rand(2,2);
julia> @ein c[i,k] := a[i,j] * b[j,k];
julia> c ≈ a * b
true
julia> c ≈ ein"ij,jk -> ik"(a,b)
trueOMEinsum.@ein_str — Macroein"ij,jk -> ik"(A,B)String macro interface which understands numpy.einsum's notation. Translates strings into StaticEinCode-structs that can be called to evaluate an einsum. To control evaluation order, use parentheses - instead of an EinCode, a NestedEinsum is returned which evaluates the expression according to parens. The valid character ranges for index-labels are a-z and α-ω.
example
julia> a, b, c = rand(10,10), rand(10,10), rand(10,1);
julia> ein"ij,jk,kl -> il"(a,b,c) ≈ ein"(ij,jk),kl -> il"(a,b,c) ≈ a * b * c
trueOMEinsum.@optein_str — Macrooptein"ij,jk,kl -> ik"(A, B, C)String macro interface that similar to @ein_str, with optimized contraction order (dimensions are assumed to be uniform).
OMEinsumContractionOrders.AbstractDecompositionType — TypeAbstractDecompositionTypeAbstract type for decomposition types, which includes TreeDecomp and PathDecomp.
OMEinsumContractionOrders.AbstractEinsum — TypeAbstractEinsumAbstract 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.
OMEinsumContractionOrders.BipartiteResult — TypeBipartiteResult{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.
OMEinsumContractionOrders.CodeOptimizer — TypeCodeOptimizerAbstract type for code optimizers.
OMEinsumContractionOrders.CodeSimplifier — TypeCodeSimplifierAbstract type for code simplifiers.
OMEinsumContractionOrders.CodeSlicer — TypeCodeSlicerAbstract type for code slicers.
OMEinsumContractionOrders.EinCode — TypeEinCode{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
3OMEinsumContractionOrders.ExactTreewidth — Typeconst 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.
OMEinsumContractionOrders.GreedyMethod — TypeGreedyMethod{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))temperatureis 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).
OMEinsumContractionOrders.HyperND — TypeHyperND(;
dis = KaHyParND(),
algs = (MF(), AMF(), MMD()),
level = 6,
width = 120,
scale = 100,
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 parameter scale controls discretization of the index weights:
weight(i) := scale * log2(dim(i))where dim(i) is the dimension of the index i.
The line graph is partitioned using the algorithm dis. OMEinsumContractionOrders currently supports two partitioning algorithms, both of which require importing an external library.
| type | package |
|---|---|
METISND | Metis.jl |
KaHyParND | KayHyPar.jl |
The optimizer is implemented using the tree decomposition library CliqueTrees.jl.
Arguments
dis: graph partitioning algorithmalgs: tuple of elimination algorithms.level: maximum levelwidth: minimum widthimbalances: imbalance parametersscore: a function to evaluate the quality of the contraction tree. Default isScoreFunction().
OMEinsumContractionOrders.KaHyParBipartite — TypeKaHyParBipartite{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_targetis the target space complexity, defined aslog2(number of elements in the largest tensor),imbalancesis a KaHyPar parameter that controls the group sizes in hierarchical bipartition,max_group_sizeis the maximum size that allowed to used greedy search,sub_optimizeris the sub-optimizer used to find the contraction order when the group size is small enough.
References
OMEinsumContractionOrders.MergeGreedy — TypeMergeGreedy <: 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).
OMEinsumContractionOrders.MergeVectors — TypeMergeVectors <: CodeSimplifier
MergeVectors()Contraction code simplifier (in order to reduce the time of calling optimizers) that merges vectors to closest tensors.
OMEinsumContractionOrders.NestedEinsum — TypeNestedEinsum{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 nodetensorindex: 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.
OMEinsumContractionOrders.NetworkSimplifier — TypeNetworkSimplifier{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 ofNestedEinsumobjects.
OMEinsumContractionOrders.PathDecomp — TypePathDecomp <: AbstractDecompositionTypePath decomposition type.
OMEinsumContractionOrders.PathSA — MethodPathSA(; βs=0.01:0.05:15, ntrials=10, niters=50, score=ScoreFunction())Optimize the einsum contraction pattern using the simulated annealing on tensor expression tree, with path decomposition.
Fields
βs,ntrials,nitersandscoreare the same as inTreeSA.
OMEinsumContractionOrders.SABipartite — TypeSABipartite{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_targetis the target space complexity, defined aslog2(number of elements in the largest tensor),ntrialsis the number of repetition (with different random seeds),βsis a list of inverse temperature1/T,nitersis the number of iteration in each temperature,max_group_sizeis the maximum size that allowed to used greedy search,sub_optimizeris the optimizer for the bipartited sub graphs, one can chooseGreedyMethod()orTreeSA(),initializeris the partition configuration initializer, one can choose:randomor:greedy(slow but better).
References
OMEinsumContractionOrders.ScoreFunction — TypeScoreFunctionA 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 thesc_weightwill be set to 0 automatically, default is 0.0.
OMEinsumContractionOrders.SlicedEinsum — TypeSlicedEinsum{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 aNestedEinsumobject.
OMEinsumContractionOrders.TreeDecomp — TypeTreeDecomp <: AbstractDecompositionTypeTree decomposition type.
OMEinsumContractionOrders.TreeSA — TypeTreeSA{IT, DT} <: CodeOptimizer
TreeSA(; βs=collect(0.01:0.05:15), ntrials=10, niters=50, initializer=:greedy, score=ScoreFunction(), decomposition_type=TreeDeomp())Optimize the einsum contraction pattern using the simulated annealing on tensor expression tree.
Fields
ntrials,βsandnitersare annealing parameters, doingntrialsindepedent annealings, each has inverse tempteratures specified byβs, in each temperature, donitersupdates of the tree.initializerspecifies how to determine the initial configuration, it can be:greedy,:randomor:specified. If the initializer is:specified, the inputcodeshould be aNestedEinsumobject.scorespecifies the score function to evaluate the quality of the contraction tree, it is a function of time complexity, space complexity and read-write complexity.decomposition_typespecifies the type of decomposition to use, it can beTreeDeomporPathDecomp.
References
Breaking changes:
nslicesis removed, since the slicing part is now separated from the optimization part, seeslice_codefunction andTreeSASlicer.greedy_methodis removed. If you want to have detailed control of the initializer, please pre-optimize the code with another method and then use:specifiedto initialize the tree.
OMEinsumContractionOrders.TreeSASlicer — TypeTreeSASlicer{IT, LT} <: CodeSlicerA 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,βsandnitersare annealing parameters, doingntrialsindepedent annealings, each has inverse tempteratures specified byβs, in each temperature, donitersupdates 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) xoptimization_ratio.score::ScoreFunction: A function to evaluate the quality of the contraction tree. Default isScoreFunction(sc_target=30.0).decomposition_type::AbstractDecompositionType: The type of decomposition to use. Default isTreeDecomp().
References
OMEinsumContractionOrders.Treewidth — Typestruct 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:
| Algorithm | Description | Time Complexity | Space Complexity |
|---|---|---|---|
AMF | approximate minimum fill | O(mn) | O(m + n) |
MF | minimum fill | O(mn²) | - |
MMD | multiple minimum degree | O(mn²) | O(m + n) |
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)
ab, ba -> a
├─ ab
└─ bcf, acf -> ba
├─ bcef, e -> bcf
│ ├─ bcef
│ └─ e
└─ acd, df -> acf
├─ acd
└─ dfOMEinsumContractionOrders.contraction_complexity — Methodcontraction_complexity(eincode, size_dict) -> ContractionComplexityReturns the time, space and read-write complexity of the einsum contraction. The returned ContractionComplexity object contains 3 fields:
tc: time complexity defined aslog2(number of element-wise multiplications).sc: space complexity defined aslog2(size of the maximum intermediate tensor).rwc: read-write complexity defined aslog2(the number of read-write operations).
OMEinsumContractionOrders.einexpr_to_matrix! — Methodeinexpr_to_matrix!(marker, ixs, iy, size_dict)Construct the weighted incidence matrix correponding to an Einstein summation expression. Returns a quadruple (weights, ev, ve, el).
Each Einstein summation expression has a set E ⊆ L of indices, a set V := {1, …, |V|} of (inner) tensors, and an outer tensor * := |V| + 1. Each tensor v ∈ V is incident to a sequence ixs[v] of indices, and the outer tensor is incident to the sequence iy. Note that an index can appear multiple times in ixs[v], e.g.
ixs[v] = ('a', 'a', 'b').Each index l ∈ E also has a positive dimension, given by size_dict[l].
The function einexpr_to_matrix does two things. First of all, it enumerates the index set E, mapping each index to a distinct natural number.
el: {1, …, |E|} → E
le: E → {1, …, |E|}Next, it constructs a vector weights: {1, …, |E|} → [0, ∞) satisfying
weights[e] := log2(size_dict[el[e]]),and a sparse matrix ve: {1, …, |V| + 1} × {1, …, |E|} → {0, 1} satisfying
ve[v, e] := { 1 if el[e] is incident to v
{ 0 otherwiseWe can think of the pair H := (weights, ve) as an edge-weighted hypergraph with incidence matrix ve.
OMEinsumContractionOrders.embed_simplifier — Methodembed_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 aNetworkSimplifierobject.
Returns
- A new
NestedEinsumobject.
OMEinsumContractionOrders.flop — Methodflop(eincode, size_dict) -> IntReturns the number of iterations, which is different with the true floating point operations (FLOP) by a factor of 2.
OMEinsumContractionOrders.getixsv — Functiongetixsv(code::AbstractEinsum) -> Vector{Vector{LT}}Returns the input indices of the einsum notation. Each vector represents the labels associated with a input tensor.
OMEinsumContractionOrders.getiyv — Functiongetiyv(code::AbstractEinsum) -> Vector{LT}Returns the output index of the einsum notation.
OMEinsumContractionOrders.label_elimination_order — Methodlabel_elimination_order(code) -> VectorReturns 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.
OMEinsumContractionOrders.labeltype — Methodlabeltype(code::AbstractEinsum) -> TypeReturns the data type to represent the labels in the einsum notation.
OMEinsumContractionOrders.matrix_to_tree! — Methodmatrix_to_tree!(marker, weights, ev, ve, el, alg)Construct a tree decomposition of an edge-weighted hypergraph using the elimination algorithm alg. We ensure that the indices incident to the outer tensor are contained in the root bag of the tree decomposition.
OMEinsumContractionOrders.optimize_code — Methodoptimize_code(eincode, size_dict, optimizer = GreedyMethod(); slicer=nothing, simplifier=nothing, permute=true) -> optimized_eincodeOptimize the einsum contraction code and reduce the time/space complexity of tensor network contraction. Returns a NestedEinsum instance. Input arguments are
Arguments
eincodeis an einsum contraction code instance, one ofDynamicEinCode,StaticEinCodeorNestedEinsum.sizeis a dictionary of "edge label=>edge size" that contains the size information, one can useuniformsize(eincode, 2)to create a uniform size.optimizeris aCodeOptimizerinstance, should be one ofGreedyMethod,Treewidth,KaHyParBipartite,SABipartiteorTreeSA. Check their docstrings for details.
Keyword Arguments
sliceris for slicing the contraction code to reduce the space complexity, default is nothing. Currently onlyTreeSASliceris supported.simplifieris one ofMergeVectorsorMergeGreedy. Default is nothing.permuteis 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());OMEinsumContractionOrders.optimize_greedy — Methodoptimize_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.
OMEinsumContractionOrders.optimize_sa — Methodoptimize_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
OMEinsumContractionOrders.optimize_tree — Methodoptimize_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.
OMEinsumContractionOrders.optimize_treewidth — Methodoptimize_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.
OMEinsumContractionOrders.peak_memory — Methodpeak_memory(code, size_dict::Dict) -> IntEstimate peak memory in number of elements.
OMEinsumContractionOrders.readjson — Methodreadjson(filename::AbstractString)Read the contraction order from a JSON file.
Arguments
filename: the name of the file to read from.
OMEinsumContractionOrders.simplify_code — Methodsimplify_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 beMergeVectorsorMergeGreedy.
Returns
- A tuple of
(NetworkSimplifier, newcode), wherenewcodeis a newEinCodeobject.
OMEinsumContractionOrders.slice_code — Methodslice_code(code, size_dict, slicer) -> sliced_codeSlice the einsum contraction code to reduce the space complexity, returns a SlicedEinsum instance.
Arguments
codeis aNestedEinsuminstance.size_dictis a dictionary of "edge label=>edge size" that contains the size information, one can useuniformsize(eincode, 2)to create a uniform size.sliceris aCodeSlicerinstance, currently onlyTreeSASliceris supported.
OMEinsumContractionOrders.tree_greedy — Methodtree_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).
OMEinsumContractionOrders.tree_to_einexpr! — Methodtree_to_einexpr!(marker, tree, ve, el, ixs, iy)Transform a tree decomposition into a contraction tree.
OMEinsumContractionOrders.uniformsize — Methoduniformsize(code::AbstractEinsum, size::Int) -> DictReturns a dictionary that maps each label to the given size.
OMEinsumContractionOrders.uniquelabels — Methoduniquelabels(code::AbstractEinsum) -> Vector{LT}Returns the unique labels in the einsum notation. The labels are the indices of the tensors.
OMEinsumContractionOrders.viz_contraction — Methodviz_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 isStressLayout().framerate: The frame rate of the animation. Default is10.filename: The name of the output file, with.gifor.mp4extension. Default is a temporary file with.mp4extension.show_progress: Whether to show progress information. Default istrue.
Returns
- the path of the generated file.
OMEinsumContractionOrders.viz_eins — Methodviz_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: TheAbstractEinsumobject 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. Ifnothing, 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 ofGraphDisplayConfigfor more information.kwargs...: Additional keyword arguments to be passed to theGraphVizconstructor.
OMEinsumContractionOrders.writejson — Methodwritejson(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 aNestedEinsumor aSlicedEinsumobject.