Einsum, Tensor Contractions & Convolution
The real notation behind attention, batched ops, and conv-as-matmul.
On this page
Beginner: einsum ("Einstein summation") is a compact notation for describing sums, transposes, and multiplications across the axes of one or more tensors, all in a single short expression — instead of writing nested loops or memorizing which function name does which specific operation.
Intermediate: the pattern is always the same: label each tensor's axes with letters, then say which output axes survive. Repeated letters across inputs get multiplied and summed over ("contracted"); letters that appear in the output are kept. Matrix multiplication is 'ij,jk->ik'; a dot product is 'i,i->'; a batch of matrix multiplications is 'bij,bjk->bik'.
Advanced: once you're comfortable with einsum, the core computation of attention — scores = queries · keysᵀ, batched across many heads and sequence positions at once — is a single readable line instead of a tangle of reshapes and transposes, which is exactly why real Transformer implementations lean on it heavily.
The einsum string is a direct, literal transcription of the summation formula — which is the entire point of the notation.
Chaining three matrices, ABC, is associative (section 1.5) — the result is identical regardless of grouping — but the amount of work is not. Suppose A is (100×5), B is (5×100), and C is (100×5).
Computing (AB)C: first AB costs about 100×5×100 = 50,000 multiplications and produces a (100×100) matrix; then (AB)C costs 100×100×5 = 50,000 more — about 100,000 total.
Computing A(BC) instead: BC costs 5×100×5 = 2,500 and produces a tiny (5×5) matrix; then A(BC) costs 100×5×5 = 2,500 more — only 5,000 total, a 20× saving, purely from choosing a smarter grouping of the exact same product.
Where this is used: this is exactly the optimization opt_einsum (section 1.18's expert note) performs automatically for any chain of three or more tensors — searching for the contraction order with the lowest total cost before executing anything.
Walk through einsum('ij,jk->ik', A, B) by hand for A = [[1, 2], [3, 4]] and B = [[5, 6], [7, 8]]. Output entry (0,0): j is repeated (contracted), so sum over j of A[0,j]·B[j,0] = (1)(5) + (2)(7) = 19. Output entry (0,1): sum over j of A[0,j]·B[j,1] = (1)(6) + (2)(8) = 22. Repeating for every (i,k) pair reproduces exactly the ordinary matrix product A @ B — because that's precisely what this einsum string was written to mean.
That last line is, almost verbatim, the first step of every attention layer in every Transformer — this is what "attention is just dot products" looks like in real code.
import numpy as np
A = np.random.rand(3, 4)
B = np.random.rand(4, 5)
C1 = A @ B
C2 = np.einsum('ij,jk->ik', A, B)
print(np.allclose(C1, C2)) # True
# Batched matmul, e.g. one small matrix multiply per item in a batch:
batch_A = np.random.rand(8, 3, 4)
batch_B = np.random.rand(8, 4, 5)
batch_C = np.einsum('bij,bjk->bik', batch_A, batch_B)
print(batch_C.shape) # (8, 3, 5)
# Simplified scaled dot-product attention scores:
queries = np.random.rand(2, 6, 16) # (batch, seq_len, dim)
keys = np.random.rand(2, 6, 16)
scores = np.einsum('bqd,bkd->bqk', queries, keys) / np.sqrt(16)
print(scores.shape) # (2, 6, 6) -> one score per query/key pair, per batch item- Convolutional layers can be implemented as a structured matrix multiplication using the "im2col" technique — unfolding overlapping image patches into rows of a matrix, then computing the whole convolution as a single matmul. This is exactly why hardware optimized for matmul (Tensor Cores, section 1.5) is also fast at convolution.
- Toeplitz matrices (constant along each diagonal) are the precise linear-algebra object behind 1D convolution — multiplying by a Toeplitz matrix is convolving with a fixed kernel.
- Modern ML compilers (XLA, TVM, Triton) can fuse and optimize einsum-style expressions automatically, often outperforming manually written loops or even hand-tuned library calls.
- Getting an einsum string subtly wrong (e.g. mismatched or misplaced repeated letters) — it usually still runs, just produces a silently wrong shape or wrong numbers. Always check
.shapeof the result against what you expected. - Forgetting that any letter missing from the output string is summed over — a common mistake is accidentally contracting over an axis you actually wanted to keep.
Going deeper
The name comes from Einstein's own notational shortcut in general relativity and tensor calculus: when the same index appears twice in a term, summation over it is implied without writing a Σ — modern ML libraries adopted the same convention because tensor contractions are exactly as central to deep learning as they are to physics.
At the master level: einsum expressions have a well-defined but non-trivial optimal contractionorder for more than two tensors — a naive left-to-right evaluation can be asymptotically far slower than an optimally ordered one, which is why libraries like opt_einsum exist specifically to search for the cheapest contraction path before executing anything.