API

ArrowPattern{L,R}

The left and right sides of a pattern, made by @einops_str and stored as type parameters, so the pattern is known at compile time. A side is a tuple: a name is a Symbol, a group a nested tuple, and ... an Einops.Ellipsis(). When commas separate the operands of an einsum, the left side is an Einops.List of one such tuple per operand.

A pattern shows as the einops"..." string that makes it.

Examples

julia> p = einops"a b c -> c (b a)"
einops"a b c -> c (b a)"

julia> typeof(p)
Einops.ArrowPattern{(:a, :b, :c), (:c, (:b, :a))}

BoundPattern

An ArrowPattern together with the lengths of some of its axes, which @einops_str makes when values are interpolated into a pattern with $. The ops pass the lengths on as keywords, so the values stay out of the type.

Examples

julia> n = 3;

julia> einops"($n b) -> b"
einops"(3 b) -> b"

@einops_str -> Union{ArrowPattern,BoundPattern,Val}

For parity with Python implementation.

In a pattern, 1 is an axis of length one, and inside parentheses it adds nothing: (1 b) is (b). Any other integer is an axis given only by its length, which can be reduced (on the left of reduce) or created (on the right of repeat). Integers and names given a length by a keyword have that length checked against the input.

$x interpolates an axis of length x, as if the value were written there: it too can only be reduced or created, and each occurrence is an axis of its own. The value is not part of the pattern's type, so it may be computed at run time without making the pattern's type unknown; the pattern then carries it along as a BoundPattern.

Examples

julia> einops"embed token (head batch) -> (embed head) token batch"
einops"embed token (head batch) -> (embed head) token batch"

julia> typeof(einops"a 1 b ... -> (b a) ...")
Einops.ArrowPattern{(:a, 1, :b, Einops.Ellipsis()), ((:b, :a), Einops.Ellipsis())}

julia> einops"i j, j k -> i k" # for einsum
einops"i j, j k -> i k"

julia> einops"a b _ d" # for parse_shape
Val{(:a, :b, -, :d)}()

julia> einops"i j * k" # for pack/unpack
Val{(:i, :j, *, :k)}()

julia> k = 2;

julia> reduce(sum, rand(4, 6), einops"(h $k) (w $k) -> h w") |> size
(2, 3)

parse_shape(x, pattern)

Capture the shape of an array in a pattern by naming dimensions, with _ to ignore a dimension and ... to ignore any number of them. Without an ellipsis, a tuple of Symbols and - works too.

Duplicate names are allowed, but they must refer to dimensions of the same size.

Examples

julia> parse_shape(rand(2,3,4), (:a, :b, -))
(a = 2, b = 3)

julia> parse_shape(rand(2,3), (-, -))
NamedTuple()

julia> parse_shape(rand(2,3,4,5), einops"first second third fourth")
(first = 2, second = 3, third = 4, fourth = 5)

julia> parse_shape(rand(2,3,4), einops"a b ...")
(a = 2, b = 3)

julia> parse_shape(rand(2,2,4), einops"a a b")  # duplicate 'a' with same size
(a = 2, b = 4)

@rearrange(x, "pattern", context...)

Shorthand for rearrange that takes the pattern as a string literal instead of einops"...". The reshape/permute plan is inlined at the call site (no @generated function, no Base.Pairs keyword object) whenever the pattern can be sized at macro-expansion time, otherwise it falls back to a rearrange call.

Examples

julia> x = rand(2, 3, 5);

julia> @rearrange(x, "a b c -> c b a") == rearrange(x, einops"a b c -> c b a")
true

rearrange(x, pattern::ArrowPattern; context...)

Rearrange the axes of x according to pattern.

Can always be expressed as a reshape + permutedims + reshape.

Examples

julia> x = rand(2,3,5);

julia> y = rearrange(x, einops"a b c -> c b a");

julia> size(y)
(5, 3, 2)

julia> y == permutedims(x, (3,2,1))
true

julia> z = rearrange(x, einops"a b c -> a (c b)");

julia> size(z)
(2, 15)

julia> z == reshape(permutedims(x, (1,3,2)), 2,5*3)
true

@reshape(x, "pattern", context...)

Shorthand for reshape that takes the pattern as a string literal instead of einops"...". See @rearrange for inlining behaviour.

Examples

julia> x = rand(2, 3);

julia> @reshape(x, "a b -> (a b)") == reshape(x, einops"a b -> (a b)")
true

reshape(x::AbstractArray, pattern::ArrowPattern; context...)

Reshape x according to pattern.

Unlike rearrange, this does not permute dimensions - symbols must appear in the same order on both sides.

Examples

julia> x = rand(1, 6, 2, 3);

julia> reshape(x, einops"1 (a b) ... -> a (b ...)"; a=2) |> size
(2, 18)

julia> x = rand(2, 3, 4);

julia> reshape(x, einops"a b c -> (a b) c") |> size
(6, 4)

@reduce(f, x, "pattern", context...)

Shorthand for reduce that takes the pattern as a string literal instead of einops"...". See @rearrange for inlining behaviour.

A value interpolated with $ is an axis of that length, as in @einops_str, and the macro passes it on as a keyword.

Examples

julia> x = rand(2, 3);

julia> @reduce(sum, x, "a b -> a") == reduce(sum, x, einops"a b -> a")
true

julia> k = 2;

julia> @reduce(sum, rand(4, 6), "(h $k) (w $k) -> h w") |> size
(2, 3)

reduce(f::Function, x::AbstractArray, pattern::ArrowPattern; context...)

Reduce an array over the dimensions specified by the pattern, using e.g. sum, prod, minimum, maximum, any, all, or Statistics.mean.

f must accept a dims::Tuple{Vararg{Int}} keyword argument, allowing for reduction over specific dimensions. This should reduce the specified dimensions to singletons, but not drop them.

Examples

julia> x = randn(64, 32, 35);

julia> y = reduce(sum, x, einops"c b t -> c b");

julia> size(y)
(64, 32)

julia> y == dropdims(sum(x, dims=3), dims=3)
true

julia> using Statistics: mean

julia> z = reduce(mean, x, einops"c b (t5 t) -> (t5 c) b", t5=5);

julia> size(z)
(320, 32)

julia> z == reshape(permutedims(dropdims(mean(reshape(x, 64,32,5,7), dims=4), dims=4), (3,1,2)), 320,32)
true

@repeat(x, "pattern", context...)

Shorthand for repeat that takes the pattern as a string literal instead of einops"...". See @rearrange for inlining behaviour.

A value interpolated with $ is an axis of that length, as in @einops_str, and the macro passes it on as a keyword.

Examples

julia> x = rand(2, 3);

julia> @repeat(x, "a b -> a b c", c=4) == repeat(x, einops"a b -> a b c", c=4)
true

julia> k = 2;

julia> @repeat(x, "a b -> ($k a) ($k b)") |> size
(4, 6)

repeat(x::AbstractArray, pattern::ArrowPattern; context...)

Repeat elements of x along specified axes.

Examples

julia> x = rand(2,3);

julia> y = repeat(x, einops"a b -> a b 1 r", r=2);

julia> size(y)
(2, 3, 1, 2)

julia> y == reshape(repeat(x, 1,1,2), 2,3,1,2)
true

julia> z = repeat(x, einops"a b -> a (b r)", r=2);

julia> size(z)
(2, 6)

julia> z == reshape(repeat(x, 1,1,2), 2,6)
true

@einsum(arrays..., "pattern", context...)

Shorthand for einsum that takes the pattern as a string literal instead of einops"...". Always expands to an einsum call (it is not type stable and routes through OMEinsum rather than a reshape/permute plan).

Examples

julia> x, y = rand(2, 3), rand(3, 4);

julia> @einsum(x, y, "i j, j k -> i k") == einsum(x, y, einops"i j, j k -> i k")
true

einsum(arrays..., pattern; optimizer=OMEinsum.GreedyMethod())

Compute the einsum operation specified by the pattern.

This function supports rearrange-style grouped axes, and can in some cases automatically resolve size ambiguities. Note that upstream Python einops.einsum does not yet do this (as of 0.8.1).

Commas in patterns separate operands, including empty patterns for scalar arrays.

Examples

julia> x, y = rand(2,3), rand(3,4);

julia> einsum(x, y, einops"i j, j k -> i k") == x * y
true

contract(ixs, iout, xs...; optimizer=OMEinsum.GreedyMethod())

Backend entry point of einsum: xs[k] carries the flat symbol labels ixs[k], and the result carries the labels iout. Ellipses and grouped axes are already resolved by the time this is called. One or two operands contract through OMEinsum; more contract pairwise, in the order optimizer finds, each pair through the two-operand method. So an array type takes over every einsum by adding methods for one and two operands, as the Reactant extension does for two.

Examples

julia> x, y = rand(2, 3), rand(3, 4);

julia> Einops.contract(((:i, :j), (:j, :k)), (:i, :k), x, y) ≈ x * y
true

pack(unpacked_arrays, pattern)

Pack a vector of arrays into a single array according to the pattern.

Examples

julia> inputs = [rand(2,3,5), rand(2,3,7,5), rand(2,3,7,9,5)];

julia> packed_array, packed_shapes = pack(inputs, (:i, :j, *, :k));

julia> size(packed_array)
(2, 3, 71, 5)

julia> packed_shapes
3-element Vector{NTuple{N, Int64} where N}:
 ()
 (7,)
 (7, 9)

unpack(packed_array, packed_shapes, pattern)

Unpack a single array into a vector of arrays according to the pattern.

Examples

julia> inputs = [rand(2,3,5), rand(2,3,7,5), rand(2,3,7,9,5)];

julia> inputs == unpack(pack(inputs, (:i, :j, *, :k))..., (:i, :j, *, :k))
true

julia> packed_array = rand(2,3,16);

julia> packed_shapes = [(), (7,), (4, 2)];

julia> unpack(packed_array, packed_shapes, einops"i j *") .|> size
3-element Vector{Tuple{Int64, Int64, Vararg{Int64}}}:
 (2, 3)
 (2, 3, 7)
 (2, 3, 4, 2)