einsum - From Index Notation to Code

Hello all. This is a short one, and it’s really a follow-on to the MLP forward and backward pass post. In that post every gradient was derived twice: once in index form, with explicit sums over explicit indices, and once in packed matrix form with transposes and outer products. The index form was there because it’s unambiguous and there’s nothing to get wrong; the packed form was there because that’s what you type into numpy.

It turns out you never needed the second step. The index-form equations are the code, in a notation called einsum. The translation is mechanical, takes about ten seconds per equation, and removes every decision about which argument to transpose.

This post covers the rule, the translation, the vocabulary, and the one property that makes it worth switching to. Everything from here on in this series will use it, so it’s worth the twenty minutes.

The index letters are the same as in the MLP post: \(i\) for the input dimension, \(j\) for hidden, \(k\) for output. Code is PyTorch, though np.einsum takes byte-identical strings if you prefer numpy.

Let’s get started.

The rule

An einsum call takes a subscript string and some tensors:

torch.einsum("ji,i->j", W, x)

The string has two halves separated by ->.

Left of the arrow: one group of letters per input tensor, comma separated. Each letter names one axis, in order. So ji says “W has two axes; call the first one j and the second one i”, and i says “x has one axis, and it’s the same i”.

Right of the arrow: the letters of the output, in the order you want them.

And then two rules, which are the whole notation:

1. A letter that appears on the left but not on the right is summed over. These are the contracted indices.
2. A letter that appears on the right is kept. These are the free indices.

So "ji,i->j" reads: i appears on the left but not the right, so sum over it; j survives to the output. Which is:

That is the definition of a matrix-vector product, and it is also, letter for letter, equation (1) from the MLP post.

The thing worth noticing is what you don’t write. There is no \(\sum\) symbol anywhere in the string. You never say “sum over \(i\)”; you say “the output has a \(j\) and no \(i\)”, and the summation is inferred from that. You specify the shape of the result, and the contraction follows. That inversion is the entire mental shift, and everything else is bookkeeping.

The translation recipe

Given any equation in index form:

1. Delete the \(\sum\) signs.
2. For each tensor on the right-hand side, write its indices in order. Comma separate them.
3. Write ->.
4. Write the indices of the left-hand side.

That’s it. Four steps, no thinking required, and no point at which you have to reason about transposes or argument order.

Reading the string off the math

Let’s do it with the actual equations from the MLP post rather than toy ones.

Hidden pre-activation, equation (1):

The sum has two tensors, \(W\) with indices \(j, i\) and \(x\) with index \(i\). The left-hand side has index \(j\). So:

ha = torch.einsum("ji,i->j", W, x) + b

The bias is added outside, since it isn’t part of the contraction.

Output pre-activation, equation (3), identical in shape:

Output weight gradient, equation (15):

There is no sum here at all. Both indices survive to the left-hand side, so both appear in the output:

dLdV = torch.einsum("k,j->kj", dLdya, h)

No letter is contracted, so nothing is summed, and what you get is an outer product. Worth pausing on: einsum didn’t need a different function for this. In numpy the packed version needed np.outer, a different name from the @ used two lines earlier, because @ between two 1-D arrays would have given a dot product instead. Here it’s the same notation with different letters.

Hidden activation gradient, equation (17), the one with the transpose:

Two tensors: \(V\) with indices \(k, j\), and the incoming gradient with index \(k\). Output index \(j\), so \(k\) is summed:

dLdh = torch.einsum("kj,k->j", V, dLdya)

No .T anywhere. We’ll come back to this one, because it’s the whole argument.

Operand order is yours to pick, incidentally, as long as the string matches it. torch.einsum("k,kj->j", dLdya, V) is the identical computation with the arguments swapped. The only thing that has to line up is each letter group with the tensor it describes.

Input weight gradient, equation (19):

Five equations, five strings, and at no point did we think about shapes, argument order, or transposes. We read letters off the page.

The vocabulary

Almost everything you need is one of these. Every string below is verified.

Two pairs in that table are worth staring at, because between them they teach the whole notation.

"i,i->i" against "i,i->". Identical inputs. One letter of difference in the output. The first keeps \(i\), so nothing is summed and you get an elementwise product, shape (5,). The second drops \(i\), so it’s summed away and you get a dot product, shape (), a scalar. Two of the most commonly confused operations in array code, and the notation makes the difference visible rather than something you infer from which function name was called.

"ij->ji" against "ii->i". The first is a transpose: one tensor, no multiplication, just relabeling axes. The second has a letter repeated within a single operand, which means “walk the diagonal”, giving (4,) from a (4,4) matrix. A repeated letter across two operands aligns them for multiplication; a repeated letter inside one operand takes its diagonal. Different meanings, and the second one surprises people.

einsum never transposes for you

This is the property that makes the notation worth adopting, and it’s best framed as the question I had when I first used it.

When computing \(\dfrac{\partial L}{\partial h} = V^\top \dfrac{\partial L}{\partial y_a}\) with einsum, does it transpose \(V\) implicitly? Should I write "jk" instead of "kj" to represent the transpose?

No, and no. einsum never transposes implicitly, and it never needs to.

Here is why the question dissolves. A transpose in a matrix product is not really an operation you want; it’s a workaround. Matrix multiplication only ever contracts the second index of the left operand against the first index of the right one. When the index you actually want to sum over is in the wrong position, you transpose to move it there. The transpose is you rearranging the data to fit the notation’s fixed contraction rule.

einsum has no fixed contraction rule. You name the index you want summed. So there is nothing to work around.

Concretely: \(V\) has shape \(N_o \times N_h\), so its entries are \(V_{kj}\). The derivation gave

The sum runs over \(k\), so \(k\) is the letter absent from the output:

dLdh = torch.einsum("kj,k->j", V, dLdya)

Writing "jk" for \(V\) would be a factual error about the data, not a way of expressing a transpose. It would claim \(V\)’s first axis is the hidden dimension, and \(V\)’s first axis is the output dimension. The subscripts describe the tensor’s actual layout; the output spec describes what you want. Given both, the contraction is determined, and getting a transpose “right” stops being a decision you can make wrongly.

There’s a safety argument here too, and it’s the practical reason I switched. With @, a wrong transpose is often still a shape-valid operation. If \(N_h\) and \(N_o\) happen to be equal, or your matrices are square, or you’re in the batch dimension of something, A @ B and A.T @ B both run, both produce correctly-shaped output, and one of them is silently wrong. Your loss goes down more slowly than it should and nothing ever raises. With einsum the letters are written down next to each other, so the same mistake is either a shape error or plainly visible on the page.

The MLP, converted

Here is the whole thing side by side. Left column is the packed numpy from the MLP post, right column is the same computation as einsum.

Look at the left column. Three different spellings, @, np.outer, and .T @, for what is conceptually one operation: multiply some things together and sum over some indices. Which spelling you need depends on the ranks of the operands and on which axis you’re contracting, and you have to work that out each time. The right column is one spelling, and the letters are copied off the equation.

The full version, with the hand-derived gradients checked against autograd:

import torch

torch.manual_seed(0)
torch.set_default_dtype(torch.float64)

Ni, Nh, No = 784, 500, 10

W = torch.randn(Nh, Ni) * 0.01
b = torch.zeros(Nh)
V = torch.randn(No, Nh) * 0.01
c = torch.zeros(No)
for p in (W, b, V, c):
    p.requires_grad_(True)

x = torch.randn(Ni)                      # one training sample
t = torch.zeros(No); t[3] = 1.0          # one-hot target, class 3

# ---- forward ----
ha = torch.einsum("ji,i->j", W, x) + b   # (Nh,)
h  = torch.sigmoid(ha)                   # (Nh,)
ya = torch.einsum("kj,j->k", V, h) + c   # (No,)
y  = torch.softmax(ya, dim=0)            # (No,)
L  = -(t * torch.log(y)).sum()           # scalar

L.backward()                             # autograd, for comparison only

# ---- backward, by hand ----
with torch.no_grad():
    dLdya = y - t                                  # (No,)
    dLdV  = torch.einsum("k,j->kj", dLdya, h)      # (No, Nh)
    dLdc  = dLdya                                  # (No,)
    dLdh  = torch.einsum("kj,k->j", V, dLdya)      # (Nh,)
    dLdha = dLdh * h * (1 - h)                     # (Nh,)
    dLdW  = torch.einsum("j,i->ji", dLdha, x)      # (Nh, Ni)
    dLdb  = dLdha                                  # (Nh,)

for name, ours, ref in [("W", dLdW, W.grad), ("b", dLdb, b.grad),
                        ("V", dLdV, V.grad), ("c", dLdc, c.grad)]:
    assert ours.shape == ref.shape
    print(f"dL/d{name}: max|ours - autograd| = "
          f"{(ours - ref).abs().max().item():.2e}")

dL/dW: max|ours - autograd| = 3.47e-18
dL/db: max|ours - autograd| = 8.67e-19
dL/dV: max|ours - autograd| = 0.00e+00
dL/dc: max|ours - autograd| = 0.00e+00

Exact to float64 roundoff, and two of them bit-identical. Checking against autograd rather than finite differences is the better move when you have it: no \(\varepsilon\) to tune and no truncation error, so a real disagreement shows up as a large number instead of hiding behind approximation noise.

Note that dLdha = dLdh * h * (1 - h) stayed as a plain *. It could be written "j,j,j->j", but there’s no reason to; einsum is for contractions, and an elementwise product isn’t one. Use it where indices are being summed, and use ordinary operators where they aren’t.

Gotchas

A few things that will bite once.

Always write the arrow. einsum has an implicit mode where you omit ->, as in "ij,jk". It then sums any repeated index and returns the free indices in alphabetical order. This is fine until the alphabetical order isn’t the order you wanted, at which point you get a silently transposed result. Writing the output explicitly costs four characters and removes the failure mode entirely.

A repeated letter within one operand is a diagonal. "ii->i" extracts the diagonal; it does not square anything. If you meant an elementwise square of a matrix, that’s A * A, not einsum.

Use ... for leading dimensions you don’t care about. "...ij,...jk->...ik" is a matmul over the last two axes with any number of batch dimensions in front. Handy when you want one function to work for both batched and unbatched input.

Contraction order matters for three or more operands, and the default may not be smart. With two operands, torch.einsum lowers to matmul or bmm and you get BLAS speed. With three or more, the order in which you contract them can change the FLOP count by orders of magnitude. torch.einsum uses opt_einsum’s path optimization when it’s available; np.einsum does not optimize unless you pass optimize=True, and the naive path can be dramatically slower. If you’re contracting three or more tensors in a hot loop, check what path you’re getting.

It isn’t a universal tool. einsum expresses products and sums over aligned indices. It won’t broadcast mismatched dimensions, apply nonlinearities, or do anything with strided or windowed access. Reach for it for contractions and use normal operations for the rest.

Why this matters more as the tensors get bigger

For a two-layer MLP, honestly, @ was fine. The case for einsum gets much stronger the moment you have more than two indices, which is immediately once you reach attention.

Attention has four: batch \(b\), head \(h\), sequence position, and head dimension \(d\). And the sequence dimension appears twice, once for the query position and once for the key position, which are different axes that happen to have the same length. Call them \(s\) and \(t\).

S = torch.einsum("bhsd,bhtd->bhst", Q, K)
P = torch.softmax(S / D**0.5, dim=-1)
O = torch.einsum("bhst,bhtd->bhsd", P, V)

The packed equivalents are Q @ K.transpose(-2, -1) and P @ V, which do work, but you have to hold in your head that the transpose is over the last two axes specifically, that the batch and head dimensions are broadcasting along, and which of the two sequence axes you’re contracting. In the einsum version, the fact that we sum over \(d\) for the scores and over \(t\) for the output is stated in the string. And critically, packed matrix notation has no way to name the two sequence axes differently, so it can’t tell you which one is which; einsum gives them separate letters and the ambiguity disappears.

That is the real argument. einsum is not a shortcut for matrix multiply. It’s the notation that keeps working when the number of indices exceeds what matrix notation can address, and it lets you carry the derivation straight into the code without a translation step where errors get introduced.

Wrapping up

The summary is short. Write your equation in index form. Delete the sums. List each tensor’s indices, comma separated, then an arrow, then the indices you want out. Letters that don’t survive to the output get summed.

Doing that removes the packed-matrix step from the pipeline entirely, and with it every transpose decision, every choice between @ and outer, and the class of bugs where a wrong transpose is still shape-valid and therefore silent.

The rest of this series uses einsum throughout, starting with LayerNorm and RMSNorm next.

If you find a mistake anywhere in here, please let me know and I’ll fix it.