Understand

einsum in practice

Reading and writing index strings — every common contraction in one line, with results from real runs.

einsum(equation, t₁, t₂, …) multiplies tensors entry by entry wherever their indices agree and sums over every index that does not appear in the output. It is Einstein’s convention (see Einstein’s summation) spelled as a string.

Reading an equation

"ij,jk->ik"
 ││ ││  ││
 ││ ││  └┴─ output: the free indices, in this order
 ││ └┴───── operand 2 has indices j, k
 └┴──────── operand 1 has indices i, j
  1. Each comma-separated group names the axes of one operand, in order. The number of letters must equal that operand’s rank.
  2. The same letter in two places means the same position: those axes are aligned and must have the same extent.
  3. Letters after -> are kept, in that order. Every other letter is summed away.
  4. An empty output ("i,i->") sums everything: the result is a scalar (a rank-0 tensor).

So "ij,jk->ik" is \(C_{ik} = \sum_j A_{ij} B_{jk}\) — matrix multiplication.

The common contractions

All of these were compiled from the .hb below and run; the outputs are copied from that run.

fn dotOf(u, v)       ↦ einsum("i,i->", u, v)
fn outerOf(u, v)     ↦ einsum("i,j->ij", u, v)
fn applyOf(a, x)     ↦ einsum("ij,j->i", a, x)
fn productOf(a, b)   ↦ einsum("ij,jk->ik", a, b)
fn transposeOf(a)    ↦ einsum("ij->ji", a)
fn columnSums(a)     ↦ einsum("ij->j", a)
fn rowSums(a)        ↦ einsum("ij->i", a)
fn gram(x)           ↦ einsum("ni,nj->ij", x, x)
fn batched(a, b)     ↦ einsum("bij,bjk->bik", a, b)
fn bilinear(u, a, v) ↦ einsum("i,ij,j->", u, a, v)
fn meanOf(x)         ↦ einsum("i->", x) / rows(x)

With u = [1, 2, 3], v = [4, 5, 6], A = [[1, 2], [3, 4]] and B = [[5, 6], [7, 8]]:

Name Equation Mathematics Result
dot product "i,i->" sum of u_i v_i 32
outer product of [1, 2] and [3, 4, 5] "i,j->ij" u_i v_j [[3, 4, 5], [6, 8, 10]]
matrix × vector [1, 1] "ij,j->i" sum over j of A_ij x_j [3, 7]
matrix product "ij,jk->ik" sum over j of A_ij B_jk [[19, 22], [43, 50]]
transpose "ij->ji" A_ji [[1, 3], [2, 4]]
column sums "ij->j" sum over i of A_ij [4, 6]
row sums "ij->i" sum over j of A_ij [3, 7]
Gram matrix of [[1, 0], [1, 1], [0, 2]] "ni,nj->ij" XᵀX [[2, 1], [1, 5]]
batched product "bij,bjk->bik" one matrix product per batch entry see below
bilinear form [1, 2]ᵀ A [3, 4] "i,ij,j->" uᵀAv 61
mean of [2, 4, 9] "i->" then / rows(x) total over count 5

The batched product multiplies two stacks of 2×2 matrices pairwise. For the stacks [[[1,2],[3,4]], [[1,0],[0,1]]] and [[[5,6],[7,8]], [[9,8],[7,6]]] the result is [[[19, 22], [43, 50]], [[9, 8], [7, 6]]] — the ordinary product, then the identity leaving the second matrix unchanged.

Calling them from C# (see Calling from C# for why rank-changing formulas take Tensor arguments):

var a = new double[,] { { 1, 2 }, { 3, 4 } };
var b = new double[,] { { 5, 6 }, { 7, 8 } };

Tensor dot   = Contractions.DotOf(new[] { 1d, 2d, 3d }, new[] { 4d, 5d, 6d });        // 32
Matrix prod  = Contractions.ProductOf(a, b);                                          // [[19,22],[43,50]]
Tensor outer = Contractions.OuterOf((Tensor)new[] { 1d, 2d }, (Tensor)new[] { 3d, 4d, 5d });
Tensor sums  = Contractions.ColumnSums((Tensor)a);                                    // [4, 6]

Designing an equation

Start from the mathematics with explicit indices, then copy the letters:

  1. Write the formula with all its sums: \(\Sigma_{ij} = \frac{1}{n} \sum_{n} (x_{ni} - \mu_i)(x_{nj} - \mu_j)\).
  2. List each operand’s indices: the centred data twice, ni and nj.
  3. The output keeps the indices not summed: ij.
  4. Write it: einsum("ni,nj->ij", c, c) / rows(c).

Letters are arbitrary — "ab,bc->ac" is the same as "ij,jk->ik" — so choose ones that read well: n for observations, p for features, a for actions, s for states, t for time.

Beyond sums: masks, reductions and positions

A contraction only ever multiplies and sums. Three kinds of question are outside it, and each has its own tool:

  • Minimum and maximum are not sums. Use the intrinsics minimum(x) (all axes) or minimum(x, axis) (one axis), and likewise maximum.
  • Positions — “the diagonal”, “the upper triangle”, “the k-th column” — are not sums either. The intrinsic triangle(x, k) keeps the upper triangle from diagonal k, and from it the Prelude builds an identity matrix and an index vector 0, 1, …, n−1; “the k-th column” is then the mask index(a) = k.
  • Selection is multiplication by a 0/1 mask, as in What is a tensor?.

For example, the trace written with a mask rather than a repeated index:

def eye(a) ≔ triangle(a^0, 0) − triangle(a^0, 1)
fn traceByMask(a) ↦ einsum("ij,ij->", a, eye(a))       -- [[1, 2], [3, 4]] gives 5

Repeated indices

Einstein’s own convention writes the trace as \(A_{ii}\), and ONNX Runtime accepts the equation "ii->": run on ONNX Runtime, einsum("ii->", [[1, 2], [3, 4]]) returns 5. The host interpreter, which runs small inputs by default, implements a subset of einsum that excludes an index repeated inside one operand and refuses it with a NotSupportedException naming the equation. Prefer the mask form above (it is what the Prelude’s diag does), or pin Programs.Executor = Executor.Onnx.

The host subset also requires an explicit -> and does not accept ... (ellipsis) — write every axis — and it requires an index shared between operands to have the same extent in each (it does not broadcast an axis of extent 1). Formulas that stay inside this subset give the same answer on both executors.

Where to find more

The Linear algebra module defines transpose, mul, apply, outer, dot, diag, identity and the decompositions built from them; Statistics and Kernels show einsum over data matrices; Networks writes layers and attention as contractions.