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
true
OMEinsum.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 EinArray
s. 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
— TypeIndexGroup
Leaf 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} <: AbstractEinsum
The abstract type for contraction trees. It has two subtypes, DynamicNestedEinsum
and StaticNestedEinsum
.
OMEinsum.NestedEinsumConstructor
— TypeNestedEinsumConstructor
describes a (potentially) nested einsum. Important fields:
args
, vector of all inputs, eitherIndexGroup
objects corresponding to tensors orNestedEinsumConstructor
iy
, indices of output
OMEinsum.SlicedEinsum
— TypeSlicedEinsum{LT, Ein} <: AbstractEinsum
A 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))
false
OMEinsum.analyze_binary
— MethodGet the expected labels.
OMEinsum.asarray
— Methodasarray(x[, parent::AbstractArray]) -> AbstactArray
Return 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 aNestedEinsum
or 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 aNestedEinsum
or 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 aNestedEinsum
or 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) -> EinArray
Constructor 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
true
OMEinsum.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))
true
OMEinsum.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
: Scalex
bysx
.sy
: Scaley
bysy
.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 i
th 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)
true
OMEinsum.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
ixs
of the input not appearing in the outputiy
- a tuple of tuples of locations of an index-label in the
ixs
in a list of all index-labels - a tuple of locations of index-labels in
iy
in 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 EinIndexer
s 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))
false
e.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
Macro 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
OMEinsum.@ein
— Macro@ein A[i,k] := B[i,j] * C[j,k] # A = B * C
Macro 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)
true
OMEinsum.@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
true
OMEinsum.@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.AbstractEinsum
— TypeAbstractEinsum
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.
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
— TypeCodeOptimizer
Abstract type for code optimizers.
OMEinsumContractionOrders.CodeSimplifier
— TypeCodeSimplifier
Abstract type for code simplifiers.
OMEinsumContractionOrders.CodeSlicer
— TypeCodeSlicer
Abstract 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
3
OMEinsumContractionOrders.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))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.HyperND
— TypeHyperND(;
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.
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_target
is the target space complexity, defined aslog2(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
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 ofNestedEinsum
objects.
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_target
is the target space complexity, defined aslog2(number of elements in the largest tensor)
,ntrials
is the number of repetition (with different random seeds),βs
is a list of inverse temperature1/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 chooseGreedyMethod()
orTreeSA()
,initializer
is the partition configuration initializer, one can choose:random
or:greedy
(slow but better).
References
OMEinsumContractionOrders.ScoreFunction
— TypeScoreFunction
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 thesc_weight
will 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 aNestedEinsum
object.
OMEinsumContractionOrders.TreeSA
— TypeTreeSA{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
andniters
are annealing parameters, doingntrials
indepedent annealings, each has inverse tempteratures specified byβs
, in each temperature, doniters
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 inputcode
should be aNestedEinsum
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
Breaking changes:
nslices
is removed, since the slicing part is now separated from the optimization part, seeslice_code
function andTreeSASlicer
.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.
OMEinsumContractionOrders.TreeSASlicer
— TypeTreeSASlicer{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
andniters
are annealing parameters, doingntrials
indepedent annealings, each has inverse tempteratures specified byβs
, in each temperature, doniters
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) xoptimization_ratio
.score::ScoreFunction
: A function to evaluate the quality of the contraction tree. Default isScoreFunction(sc_target=30.0)
.
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 |
---|---|---|---|
BFS | breadth-first search | O(m + n) | O(n) |
MCS | maximum cardinality search | O(m + n) | O(n) |
LexBFS | lexicographic breadth-first search | O(m + n) | O(m + n) |
RCMMD | reverse Cuthill-Mckee (minimum degree) | O(m + n) | O(m + n) |
RCMGL | reverse Cuthill-Mckee (George-Liu) | O(m + n) | O(m + n) |
MCSM | maximum cardinality search (minimal) | O(mn) | O(n) |
LexM | lexicographic breadth-first search (minimal) | O(mn) | O(n) |
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)
ba, ab -> a
├─ bcf, fac -> ba
│ ├─ e, bcef -> bcf
│ │ ├─ e
│ │ └─ bcef
│ └─ df, acd -> fac
│ ├─ df
│ └─ acd
└─ ab
OMEinsumContractionOrders.contraction_complexity
— Methodcontraction_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 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.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 aNetworkSimplifier
object.
Returns
- A new
NestedEinsum
object.
OMEinsumContractionOrders.flop
— Methodflop(eincode, size_dict) -> Int
Returns 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) -> 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.
OMEinsumContractionOrders.labeltype
— Methodlabeltype(code::AbstractEinsum) -> Type
Returns the data type to represent the labels in the einsum notation.
OMEinsumContractionOrders.optimize_code
— Methodoptimize_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 ofDynamicEinCode
,StaticEinCode
orNestedEinsum
.size
is a dictionary of "edge label=>edge size" that contains the size information, one can useuniformsize(eincode, 2)
to create a uniform size.optimizer
is aCodeOptimizer
instance, should be one ofGreedyMethod
,Treewidth
,KaHyParBipartite
,SABipartite
orTreeSA
. Check their docstrings for details.
Keyword Arguments
slicer
is for slicing the contraction code to reduce the space complexity, default is nothing. Currently onlyTreeSASlicer
is supported.simplifier
is one ofMergeVectors
orMergeGreedy
. 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());
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) -> Int
Estimate 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 beMergeVectors
orMergeGreedy
.
Returns
- A tuple of
(NetworkSimplifier, newcode)
, wherenewcode
is a newEinCode
object.
OMEinsumContractionOrders.slice_code
— Methodslice_code(code, size_dict, slicer) -> sliced_code
Slice the einsum contraction code to reduce the space complexity, returns a SlicedEinsum
instance.
Arguments
code
is aNestedEinsum
instance.size_dict
is a dictionary of "edge label=>edge size" that contains the size information, one can useuniformsize(eincode, 2)
to create a uniform size.slicer
is aCodeSlicer
instance, currently onlyTreeSASlicer
is 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.uniformsize
— Methoduniformsize(code::AbstractEinsum, size::Int) -> Dict
Returns 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.gif
or.mp4
extension. Default is a temporary file with.mp4
extension.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
: TheAbstractEinsum
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. 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 ofGraphDisplayConfig
for more information.kwargs...
: Additional keyword arguments to be passed to theGraphViz
constructor.
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 aNestedEinsum
or aSlicedEinsum
object.