import numpy as np
v = np.array([1.0, 2.0, 3.0])
w = np.array([4.0, 5.0, 6.0])
total = 0.0
for j in range(3): # j lives only inside this loop
total += v[j] * w[j]
print(total) # the answer has no j in it32.0einsum, one index at a time
machine-learning
Every PyTorch codebase I read eventually hits a line like this:
attn = torch.einsum('bhqd,bhkd->bhqk', q, k)I could always decode these strings mechanically. Match the letters, remember that repeated letters get summed, move on. But decoding is not the same as reading. I wanted 'ij,ja->ia' to be a sentence I understand, not an incantation I pattern-match.
It turned out the whole system rests on one rule, and everything else is that rule applied to progressively bigger objects. This post builds it up in that order: one vector, two vectors, a matrix, a batch, three tensors, and finally the attention string above. Prerequisites: vectors, dot products, and basic matrix multiplication. Nothing else.
The tempting model: a lookup table
The way most people hold einsum in their head is as a table of memorized entries:
'ij,jk->ik'is matrix multiply'ij->ji'is transpose'ii->'is trace'i,i->'is dot product
The table has no entry for 'bhqd,bhkd->bhqk', and adding more entries doesn’t help, because the strings are not arbitrary codes. Each one is a compressed equation, and the compression scheme is a physics convention from 1916 called Einstein summation. Learn the scheme once and every string, including ones you have never seen, becomes legible.
Step 1: a dot product, and one rule
Start with the smallest thing that involves a sum. The dot product of two vectors:
\[\mathbf{v}\cdot\mathbf{w} = v_1 w_1 + v_2 w_2 + v_3 w_3 = \sum_{j=1}^{3} v_j\, w_j\]
The index \(j\) is doing the walking: it visits every slot, multiplies the paired entries, and the sigma adds them up.
Einstein’s convention is a single observation about equations like this: whenever an index appears exactly twice in a term, it is always being summed over, so the sigma carries no information.
\[\mathbf{v}\cdot\mathbf{w} = v_j\, w_j\]
That is the entire convention. An index written twice means “sum over it, silently.” Everything that follows is this rule meeting larger objects.
Step 2: dummy indices and free indices
The dropped sigma creates two kinds of index, and telling them apart is the entire skill. Both have exact counterparts in code you already write, so let’s build the distinction there.
A dummy index is a loop variable inside a sum. Here is \(v_j w_j\) as the loop it abbreviates:
Everything you know about the variable j transfers to the index \(j\). It exists only inside the computation. It visits every value. And once the loop finishes, it is gone: the answer is 32.0, not something indexed by \(j\). Its name is irrelevant too. Rename j to k and it is the same computation, which is why \(v_j w_j\) and \(v_k w_k\) denote the same number. An index that appears twice is a loop variable, and spotting one should trigger the same reflex: this will be summed away and won’t exist in the output.
A free index is the opposite: it survives into the output and labels its slots. In code, it is the variable that indexes the result you are building. Take matrix-vector multiplication, \(u_i = A_{ij} v_j\):
A_small = np.arange(9.0).reshape(3, 3)
u = [sum(A_small[i, j] * v[j] for j in range(3)) # j: loop variable, consumed
for i in range(3)] # i: labels the output slot
print(u)[np.float64(8.0), np.float64(26.0), np.float64(44.0)]Two nested roles, visually separated. The inner j is a dummy again, consumed by sum. The outer i is free: the output is a list indexed by i, three numbers, one per value it takes. A free index appears exactly once in each term of the equation, and the number of free indices tells you the output’s number of dimensions before you compute anything. Zero free indices, scalar. One, vector. Two, matrix.
One consistency law binds every valid equation: the free indices on the left must exactly match the free indices on the right. In loop terms: both sides must be built by the same outer for loops. A dummy on the right is fine, it is consumed by its own sum. A free index on only one side is like writing u[i] = ... where the right side loops over m instead: the program is malformed before it runs. This law catches most index mistakes before any arithmetic happens.
The analogy has one limit worth naming. A literal loop nest fixes an execution order, but the einsum equation does not: the sums can be regrouped and reordered freely, which is exactly the freedom optimizers exploit later. The loop picture is for reading the equation, not for how the machine runs it.
From here on, whenever an equation appears, I’ll mention the roles the same way: which letter is the loop variable, which letters shape the output.
Step 3: the strings are the equations
np.einsum takes the equation you would write on paper and asks you to type it, letters only:
- input tensors are described by their indices, comma-separated
- the arrow shows the free indices that survive into the output
- any index missing after the arrow is a dummy: summed
The dot product \(v_j w_j\) over two vectors \(v\) and \(w\) has no free index, so nothing follows the arrow:
print(np.einsum('j,j->', v, w))32.0Now flip one decision. Keep both indices free instead of matching them, \(M_{ij} = v_i w_j\), and no sum happens at all. Every component of \(\mathbf{v}\) meets every component of \(\mathbf{w}\):
print(np.einsum('i,j->ij', v, w))[[ 4. 5. 6.]
[ 8. 10. 12.]
[12. 15. 18.]]This is the outer product, a 3×3 grid, and the string told you its shape before running anything: two free indices in, two dimensions out. Compare the two strings. 'j,j->' and 'i,j->ij' differ only in whether the letters match, and that single choice is the difference between collapsing to a scalar and expanding to a matrix.
Two more one-liners fall out. Transposing is just reordering free indices, no sum anywhere:
M = np.einsum('i,j->ij', v, w)
print(M.shape)(3, 3)And repeating an index within a single tensor walks its diagonal. \(A_{ii}\) means sum the entries where row equals column, the trace:
print(np.einsum('ii->', M))32.0Step 4: a matrix acts on a vector
A matrix-vector product, written per component, is a dot product between row \(i\) and the vector:
\[u_i = A_{ij}\, v_j\]
This is exactly the loop nest from Step 2, now in notation: \(j\) is the loop variable, consumed by its sum; \(i\) names the output slot. One free index, so the output is a vector. The string is a direct transcription:
A_demo = np.arange(9.0).reshape(3, 3)
dotted = np.einsum('ij,j->i', A_demo, v)
print(dotted, dotted.shape)[ 8. 26. 44.] (3,)Matrix-matrix multiplication is the same sentence with one more free index along for the ride, \(C_{ik} = A_{ij}B_{jk}\): contract \(A\)’s columns against \(B\)’s rows over the shared \(j\), keep \(i\) and \(k\) to shape the result.
B_demo = np.ones((3, 3))
matmul = np.einsum('ij,jk->ik', A_demo, B_demo)
print(matmul, matmul.shape)[[ 3. 3. 3.]
[12. 12. 12.]
[21. 21. 21.]] (3, 3)Step 5: where a real matrix comes from
So far every matrix has been handed to us. The notation gets genuinely powerful when an equation produces a matrix you didn’t know was there. Here is the example that made this click for me: lets sayP> sunlight with direction \(\hat{\mathbf{s}}\) casts a point \(\mathbf{r}\) onto the ground (the plane where the third coordinate is zero). Sliding along the ray until it hits the ground gives the shadow point
\[\mathbf{r}' = \mathbf{r} - \hat{\mathbf{s}}\,\frac{\mathbf{r}\cdot\hat{\mathbf{e}}_3}{s_3}\]
where \(\hat{\mathbf{e}}_3 = (0,0,1)\) picks out height and \(s_3\) is the sun direction’s third component. The right side is linear in \(\mathbf{r}\), so a matrix is hiding in it. Index notation digs it out in three moves, and to keep each move honest we’ll drag a dummy input through all of them: the sun at \(\hat{\mathbf{s}} = (4/13, -3/13, -12/13)\), so \(s_3 = -12/13\), and a point at \(\mathbf{r} = (6, 2, 3)\), three units above the ground.
s = np.array([4/13, -3/13, -12/13])
r = np.array([6.0, 2.0, 3.0])
e3 = np.array([0.0, 0.0, 1.0])Move 1: components. Write the \(i\)-th component of each side, with the dot product in convention form:
\[r'_i = r_i - s_i\, [\hat{\mathbf{e}}_3]_j\, r_j / s_3\]
Roles: \(j\) is the loop variable, \(i\) shapes the output, present once in every term on both sides. Well-formed.
Now watch \(j\) actually loop and vanish. The piece \([\hat{\mathbf{e}}_3]_j\, r_j\) expands to \(0\cdot 6 + 0\cdot 2 + 1\cdot 3\) i.e. a dot product between the two vectors, the basis \(\hat{\mathbf{e}}_3\) vector and the point \(\mathbf{r}\), which is exactly the height of the point. Check it:
print(sum(e3[j] * r[j] for j in range(3)))3.0It collapses to \(3\), the height of the point, with no \(j\) anywhere in the answer. The equation for the first component, \(i=1\), is then plain arithmetic: \(r'_1 = 6 - (4/13)\cdot 3/(-12/13) = 6 + 1 = 7\).
print(r[0] - s[0] * 3 / s[2])7.0So the first shadow coordinate is \(7\) before any matrix exists. The remaining moves don’t change this number; they reorganize the computation so a matrix becomes visible.
Move 2: rewrite both terms using the identity matrix. Here is where we left off:
\[r'_i = r_i - s_i\, [\hat{\mathbf{e}}_3]_j\, r_j / s_3\]
The goal of this move is simple: make both terms look alike, so the next move can factor them. Two small rewrites do it.
First, \([\hat{\mathbf{e}}_3]_j\) is the list \((0,0,1)\) exactly row 3 of the identity matrix. So we can write it as \(I_{3j}\):
I = np.eye(3)
print(I[2])[0. 0. 1.]Second, the bare \(r_i\) can be written as \(I_{ij}r_j\). Row \(i\) of the identity has a single 1 in slot \(i\), so dotting it with \(\mathbf{r}\) just returns component \(i\). Check for \(i=2\): row \((0,1,0)\) against \((6,2,3)\) gives \(2\) same for the remaining values meaning \(r_i = I_{ij} r_j\) is valid for all \(i\):
print(sum(I[1, j] * r[j] for j in range(3)))2.0Why dress up a term that was already simple? Because factoring needs both terms to end in the same \(r_j\), and \(r_i\) alone has no \(j\) to offer.
Move 3: factor out \(r_j\). With both rewrites in place, every term ends in \(r_j\), and it pulls out:
\[r'_i = I_{ij}\, r_j - \frac{s_i\, I_{3j}}{s_3}\, r_j = \left(I_{ij} - \frac{s_i\, I_{3j}}{s_3}\right) r_j\]
The bracket has two free indices, \(i\) and \(j\) so it is a matrix. This is the matrix-vector pattern \([A\mathbf{r}]_i = A_{ij}r_j\) from Step 4, with
\[A_{ij} = I_{ij} - \frac{s_i\, I_{3j}}{s_3}\]
Check with the dummy input: row \(i=1\) of the bracket is \((1, 0, 1/3)\), and dotting it with \((6,2,3)\) should reproduce the \(7\) from Move 1:
A = I - np.outer(s, I[2]) / s[2] # A_ij = I_ij - s_i I_3j / s3
print(A)
print(sum(A[0, j] * r[j] for j in range(3)))[[ 1. 0. 0.33333333]
[ 0. 1. -0.25 ]
[ 0. 0. 0. ]]
7.0Same number, now produced by a matrix row. Note np.outer(s, I[2]) is exactly the \(s_i I_{3j}\) term: two free indices, no shared letter the outer product from Step 3 in the wild.
The matrix itself is readable. \(I_{3j}\) is zero unless \(j=3\), so only column 3 differs from the identity: all the shadow action lives in one column. And at \(i=3\) the two terms cancel, so row 3 is all zeros: every output has height zero the 3D-to-2D collapse written as a dead row.
Applying the full matrix finishes the example:
print(A @ r)[7. 1.25 0. ]The shadow lands at \((7,\ 5/4,\ 0)\): shifted \(+x\) because the sun direction has positive \(x\), shifted down in \(y\) because it has negative \(y\), and flat on the ground.
One bonus check. A shadow of a shadow is itself, so \(A^2 = A\):
print(np.allclose(A @ A, A))TrueMaps with this property are called projections, and this one is oblique: unless the sun is directly overhead, points travel diagonally to the plane rather than dropping straight down.
Step 6: the batch dimension appears
Now shadow four points at once. Stack them as columns of a matrix \(R\) and write the transformation of the whole batch:
\[R'_{ia} = A_{ij}\, R_{ja}\]
Read the roles. \(j\): appears twice, the loop variable, doing the actual contraction. \(i\): free, output component. \(a\): free, and it means which vector you are on. The crucial observation is what \(a\) does not do. It never enters the sum. Nothing in the equation lets column 3 influence column 2. Batching works not by accident but because the batch index is free on both sides, and the notation makes that independence visible rather than asking you to trust it.
R = np.array([[5, -1, -3, 7],
[4, -4, 1, -2],
[9, 3, 0, 12]], dtype=float)
R_batch = np.einsum('ij,ja->ia', A, R)
print(R_batch)[[ 8. -0. -3. 11. ]
[ 1.75 -4.75 1. -5. ]
[ 0. 0. 0. 0. ]]The second check confirms it: the einsum equals four independent matrix-vector products laid side by side. This is the same mechanism by which a batch dimension in a deep learning framework broadcasts through a matmul without mixing examples.
Before scaling up further, a checkpoint. Write your answer before expanding.
NoteCheckpoint: what does
np.einsum('ij,jk,ka->ia', A, A, R) compute, and what should it equal here?
Read the string one index at a time. \(j\) appears in ij and jk a dummy, contracting the first \(A\) against the second. \(k\) appears in jk and ka another dummy, contracting that result against \(R\). What survives is \(i\) (output component) and \(a\) (which vector in the batch), so the output is ia, the same shape as one batched matmul:
\[R''_{ia} = A_{ij}\, A_{jk}\, R_{ka}\]
In words: apply \(A\) twice to the whole batch, examples still kept independent by the free \(a\).
print(np.einsum('ij,jk,ka->ia', A, A, R))[[ 8. -0. -3. 11. ]
[ 1.75 -4.75 1. -5. ]
[ 0. 0. 0. 0. ]]Step 7: reading the attention string
The opening line can now be read as a sentence:
attn = torch.einsum('bhqd,bhkd->bhqk', q, k)Hunt for the loop variable first: which letter is missing after the arrow? Only \(d\), the feature dimension. So \(d\) is the one that loops and vanishes: queries and keys are contracted along features, a dot product per pair, exactly the 'j,j->' from Step 1 wearing more clothes.
Then account for the free indices. \(b\) (batch) and \(h\) (head) appear in both inputs and the output in matching positions: they ride along untouched, exactly like \(a\) did in Step 6, two batch dimensions instead of one. \(q\) comes only from the first tensor and \(k\) only from the second, and both survive: every query meets every key, the outer-product pattern from Step 3. The output is a \(q \times k\) grid of similarity scores for each batch and head.
rng = np.random.default_rng(0)
q = rng.standard_normal((2, 3, 4, 8)) # batch=2, heads=3, queries=4, features=8
k = rng.standard_normal((2, 3, 5, 8)) # batch=2, heads=3, keys=5, features=8
scores = np.einsum('bhqd,bhkd->bhqk', q, k)
print(scores.shape)(2, 3, 4, 5)The shape check alone shows the grammar working: the string predicted (2, 3, 4, 5) before the computation ran, and the einsum matches the conventional transpose-and-matmul spelling while stating its intent far more directly.
Where the mental model stops
First, reading an einsum tells you what is computed, not how fast. With three or more tensors, as in the checkpoint above, there is a choice of which pair to contract first, and for large tensors that choice can change the cost by orders of magnitude. NumPy’s optimize=True flag and the opt_einsum library exist to search for a good contraction order. The notation is about meaning; scheduling is a separate problem.
Second, not everything is an einsum. The convention expresses sums of products over matched indices: contractions, traces, transposes, outer products, and their batched combinations. Elementwise nonlinearities, softmax, max-pooling, anything with a comparison or a division by a computed sum, live outside it. When you see einsum in a codebase, you are looking at the multilinear skeleton of the computation, with the nonlinear flesh applied elsewhere. In the attention example, the softmax and the value-weighting that follow the score computation are separate lines for exactly this reason.