Einsum and Tensor Shapes

Why Does This Exist?

Multi-head attention needs to multiply a (batch, heads, seq, dim) tensor by a (batch, heads, dim, seq) tensor, contracting only the dim axis. Written with transpose, reshape and matmul that is four or five lines, each with an axis order you have to hold in your head, and any one of them can be wrong in a way that still produces a plausible shape.

Written with einsum it is one line, and the line states its own contract:

1scores = np.einsum("bhqd,bhkd->bhqk", queries, keys)

You can read off what happens without running anything. d appears on both inputs and not the output, so it is summed away. b and h appear everywhere, so they are carried along as batch. q and k each appear once, so they form the output grid. That is an attention score matrix, and the notation is the documentation.

This is the most practically useful page in the linear algebra section, because shape bugs are the most common bugs in ML and this is the tool that makes them visible.

Think of It Like This

How It Actually Works

The notation is "input1,input2->output", with one letter per axis. Three rules:

1. A letter in the inputs but not the output is summed over (contracted).
2. A letter in both inputs and the output is matched positionally — a batch or elementwise axis.
3. Output letter order sets the output axis order, so permuting letters transposes for free.

Every familiar operation is a special case:

• "ij->ji" — transpose. One input, letters reordered.
• "ij,jk->ik" — matrix multiplication. j is contracted.
• "i,i->" — dot product. Everything summed, scalar out.
• "ii->" — trace. A repeated letter on one input takes the diagonal.
• "i,j->ij" — outer product. Nothing shared, nothing summed.
• "ij->i" — row sums. j dropped, so j is summed.
• "bij,bjk->bik" — batched matmul. b carried, j contracted.

That is the entire API. Once you see that sum, transpose, matmul, trace, outer and diagonal are all one operation with different letters, the notation stops feeling exotic.

Why it is safer than the alternatives

It states the contract. "bhqd,bhkd->bhqk" tells a reader what the axes mean. x.transpose(0,1,3,2) @ y does not.

It catches errors early. Mismatched letters produce an immediate shape error, whereas a wrong transpose often yields a valid shape and wrong numbers.

It removes intermediate allocations. np.einsum("ij,jk,kl->il", A, B, C) can plan the contraction, and with optimize=True it chooses the cheaper association order. That matters: multiplying $(100 \times 1000)(1000 \times 100)(100 \times 1000)$ left-to-right versus right-to-left differs by an order of magnitude in work.

Broadcasting is explicit. Ellipsis "...ij,...jk->...ik" handles any number of leading batch axes without you counting them.

Shape discipline that pays off regardless

Einsum is most useful alongside three habits:

• Name your axes in comments or types — # (batch, seq, heads, dim) above any nontrivial tensor.
• Assert shapes at boundaries — assert scores.shape == (b, h, q, k). Cheap, and converts a silent wrong answer into a failure at the responsible line.
• Know that reshape is not transpose. Reshape re-cuts memory order; transpose moves data. Using reshape to reorder axes succeeds, matches the element count, and scrambles the contents.

Worked example

Take $A$ of shape $(2,3)$ and $B$ of shape $(3,2)$:

Write it as "ij,jk->ik" with $i = 2$, $j = 3$, $k = 2$. The letter $j$ appears on both inputs and not the output, so it is summed — and it must match, which it does at 3. The output is $(i,k) = (2,2)$.

Compute each entry as a sum over $j$:

• $(0,0)$: $1{\cdot}1 + 2{\cdot}0 + 3{\cdot}1 = 4$
• $(0,1)$: $1{\cdot}0 + 2{\cdot}1 + 3{\cdot}1 = 5$
• $(1,0)$: $4{\cdot}1 + 5{\cdot}0 + 6{\cdot}1 = 10$
• $(1,1)$: $4{\cdot}0 + 5{\cdot}1 + 6{\cdot}1 = 11$

Now change only the output letters. "ij,jk->ki" gives the transpose, $\begin{bmatrix} 4 & 10 \\ 5 & 11\end{bmatrix}$, with no transpose call — just a different letter order.

And "ij->i" on $A$ alone sums each row: $1+2+3 = 6$ and $4+5+6 = 15$, giving $[6, 15]$. The letter $j$ is absent from the output, so it is summed. Same rule, no special case.

The diagram's "bij,bjk->bik" is this with a batch letter added: $(2,3,4)$ and $(2,4,5)$ give $(2,3,5)$, with $j = 4$ matching and vanishing.

Show Me the Code

1import numpy as np2
3A = np.array([[1, 2, 3], [4, 5, 6]])          # (2, 3) = (i, j)4B = np.array([[1, 0], [0, 1], [1, 1]])        # (3, 2) = (j, k)5
6print(np.einsum("ij,jk->ik", A, B).tolist())   # -> [[4, 5], [10, 11]]7print(np.einsum("ij,jk->ki", A, B).tolist())   # -> [[4, 10], [5, 11]]  transposed8print(np.einsum("ij->i", A).tolist())          # -> [6, 15]   row sums9print(np.einsum("ij->ji", A).shape)            # -> (3, 2)    transpose10
11S = np.array([[1, 2], [3, 4]])12print(np.einsum("ii->", S), np.einsum("i,i->", A[0], A[0]))  # -> 5 14  trace, dot13
14# Batched attention scores: contract d, keep batch and head.15q = np.random.default_rng(0).standard_normal((2, 4, 6, 8))   # (b, h, q, d)16k = np.random.default_rng(1).standard_normal((2, 4, 6, 8))   # (b, h, k, d)17print(np.einsum("bhqd,bhkd->bhqk", q, k).shape)              # -> (2, 4, 6, 6)

The matrix product, its transpose, and the row sums all match the hand calculation. The last line is the real payoff: one readable line for what would otherwise be several transposes.

Watch Out For

The Quick Version

• Notation is "in1,in2->out", one letter per axis.
• A letter absent from the output is summed. A letter in both inputs and the output is matched. Output letter order sets the axis order.
• Everything is a special case: ij->ji transpose, ij,jk->ik matmul, i,i-> dot, ii-> trace, i,j->ij outer, ij->i row sums, bij,bjk->bik batched matmul.
• It states its own contract, which makes it reviewable — the main reason to use it.
• Errors usually surface as shape mismatches rather than wrong numbers.
• Use optimize=True on three or more operands; the contraction order can differ by an order of magnitude.
• ... handles any number of leading batch axes.
• For a plain 2-D product, @ is clearer and faster. Use einsum for genuinely multi-axis work.
• Omitting a letter from the output silently sums that axis — including the batch. Assert the output shape.
• A repeated letter on one operand takes a diagonal, not an elementwise product.

What to Read Next

• Matrix Multiplication is the contraction einsum generalises, and where the cost model lives.
• Vectors, Matrices and Tensors covers the shape, axis and broadcasting rules underneath all of this.
• Dot Product is the "i,i->" case.
• Distance and Similarity Metrics is a natural batched-einsum use: all-pairs similarity in one line.
• Computational Complexity for ML prices the contraction an einsum string describes, including the one that materialises a huge intermediate.
• Definitions worth a look: Tensor, Transpose, Matrix, and Scalar.
• Related interview question: What does a tensor's shape tell you?