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.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.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.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
— Methodeinsum(code::EinCode, xs, size_dict)
einsum(rule, ixs, iy, xs, size_dict)
return the tensor that results from contracting the tensors xs
according to their indices ixs
(getixs(code)
), where all indices that do not appear in the output iy
(getiy(code)
) are summed over. The result is permuted according to out
.
ixs
- tuple of tuples of index-labels of the input-tensorsxs
iy
- tuple of index-labels of the output-tensorxs
- tuple of tensorssize_dict
- a dictionary 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
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!(::EinCode, xs, y, 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
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