einsum

→Tensor

einsum(equation: str, operands: Tensor = ())

source edit

Implementing kernel

C++einsum_opfree fn

Evaluate an Einstein-summation expression over the given tensors.

Implements the general tensor contraction operator: each input is indexed by labels declared in equation and the output retains only the labels appearing on the right-hand side. Repeated labels are summed; matching labels broadcast.

Parameters

equationstr

Einsum equation of the form "<lhs1>,<lhs2>,... -> <rhs>".

*operandsTensor= ()

Tensors corresponding to the comma-separated factors on the left-hand side.

Returns

Tensor

Tensor whose axes correspond to the labels on the right-hand side of equation.

Raises

lucid._C.engine.InvalidArgument

The equation is malformed: a subscript outside [a-zA-Z], a stray ., the wrong number of operands, or an output subscript that repeats or appears in no input.

lucid._C.engine.ShapeMismatch

An operand's rank does not match its subscripts, or a subscript's sizes disagree (within one operand, or across operands without a size of 1 to broadcast).

Notes

Matrix multiplication is the canonical example: "ij,jk->ik" realises

C_{ik} = \sum_{j} A_{ij} \, B_{jk}.

More generally, einsum expresses traces ("ii->"), batched contractions ("bij,bjk->bik"), and outer products ("i,j->ij") within one consistent notation.

A subscript repeated within one operand selects that operand's diagonal over those axes ("ii->i", "bii->bi"), and must name axes of equal size. ... stands for the remaining leading axes; they broadcast across operands, right-aligned. A subscript's size may be 1 in one operand and longer in another, which broadcasts. Without -> the output is ... followed by the subscripts that occur exactly once, in sorted order.

Examples

>>> import lucid
>>> a = lucid.randn(2, 3)
>>> b = lucid.randn(3, 4)
>>> lucid.einops.einsum("ij,jk->ik", a, b).shape
(2, 4)
>>> m = lucid.tensor([[1.0, 2.0], [3.0, 4.0]])
>>> lucid.einops.einsum("ii->i", m)
tensor([1., 4.])
>>> lucid.einops.einsum("ii->", m)
tensor(5.)

>>> import lucid >>> a = lucid.randn(2, 3) >>> b = lucid.randn(3, 4) >>> lucid.einops.einsum("ij,jk->ik", a, b).shape (2, 4) >>> m = lucid.tensor([[1.0, 2.0], [3.0, 4.0]]) >>> lucid.einops.einsum("ii->i", m) tensor([1., 4.]) >>> lucid.einops.einsum("ii->", m) tensor(5.)