EinArray{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,

EinCode{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> EinCode((('i','j'),('j','k')),('i','k'))(a, b) ≈ a * b
true

EinIndexer{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,

EinIndexer{locs}(size::Tuple)

Constructor for EinIndexer for an object of size size where locs are the locations of relevant indices in a larger tuple.

IndexSize{N,T}(k::NTuple{N,T},v::NTuple{N,Int})

struct to hold the size of indices specified by their labels. Note that while a dict would work, for the small sizes we usually have, a tuple of keys and values is much faster to construct and competitive for lookup.

asarray(x::Number[, 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.

einarray(::EinCode, xs, size_dict) -> EinArray

Constructor of EinArray from an EinCode, a tuple of tensors xs and a size_dict of type IndexSize 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(EinCode((('i','j'),('j','k')),('i','k')), (a,b), sd);

julia> dropdims(sum(ea, dims=1), dims=1) ≈ a * b
true

einsum(::EinCode{ixs, iy}, xs, size_dict) where {ixs, iy}

return the tensor that results from contracting the tensors xs according to their indices ixs, where all indices that do not appear in the output iy are summed over. The result is permuted according to out.

• ixs - tuple of tuples of index-labels of the input-tensors xs

• iy - tuple of index-labels of the output-tensor

• xs - tuple of tensors

• size_dict - IndexSize-object that maps index-labels to their sizes

example

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

loop_einsum!(::EinCode, xs, y, size_dict)

inplace-version of loop_einsum, saving the result in a preallocated tensor of correct size y.

loop_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.

@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

ein"ij,jk -> ik"(A,B)

String macro interface which understands numpy.einsum's notation. Translates strings into EinCode-structs that can be called to evaluate an einsum. To control evaluation order, use parentheses - instead of an EinCode, a NestedEinsumStable 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

IndexGroup

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).

NestedEinsum

describes a (potentially) nested einsum. Important fields:

• args, vector of all inputs, either IndexGroup objects corresponding to tensors or NestedEinsum
• iy, indices of output

apply a NestedEinsum to arguments evaluates the nested einsum

getindex(A::EinArray, inds...)

return the lazily calculated entry of A at index inds.

allunique(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

bpcheck(f, args...; η = 1e-5, verbose=false)

returns a Bool indicating whether Zygote calculates the gradient of f(args...) -> scalar correctly using the relation f(x - ηg) ≈ f(x) - η|g|² with a relative tolerance of 1e-2 and an absolute tolerance of 1e-8. If verbose=true, print f(x) - f(x - ηg)and η|g|².

example

julia> using OMEinsum: bpcheck

julia> a, b = rand(2,2), rand(2,2);

julia> bpcheck(sum ∘ ein"ij,jk -> ik", a, b)
true

check_dimensions(inds::IndexSize)

check whether all non-unique indexlabels point to the same dimension - otherwise throw an error.

einsum_grad(::EinCode{ixs, iy}, xs, 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(EinCode((('i','j'),('j','k')), ('i','k')), (a,b), sd, c, 1) ≈ c * transpose(b)
true

extract the indices of the tensor that is associated with x (if x isa IndexGroup) or results from x (if x isa NestedEinsum)

extract the tensor associated with x (if x isa IndexGroup) or evaluate and return the tensor associated with x (if x isa NestedEinsum)

filliys!(neinsum::NestedEinsum)

goes through all NestedEinsum objects in the tree and saves the correct iy in them.

get_size_dict(ixs, xs)

return the IndexSize struct 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.

indices_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 output iy
• 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.

map_prod(xs, ind, indexers)

calculate the value of an EinArray with EinIndexers indexers at location ind.

match_rule(ixs, iy)
match_rule(code::EinCode{ixs, iy})

go through all operations specified in the einsum_rules-vector and return the first T for which match_rule(T, ixs, iy) returns true.

Hadamard

Ptrace rule if all indices of one ix in ixs all appear in iy or appear twice and don't appear in iy

permutation rule

a einsum code is sum.

a einsum code is trace

a einsum code is a pairwise graph.

nopermute(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

parse_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.

stabilize(nein::NestedEinsum)

turn the nested einsum into a NestedEinsumStable which is type-stable and more performant.