Tensors and Shape Manipulation — If You Can Read einsum, You Can Read Papers

The Σ_j A_ij B_jk in the paper and the x.transpose(1,2) in the code say the same thing, and einsum is the bridge between them. Three tools — axes, broadcasting, contraction — are enough to write attention in a single line.

Where people actually get stuck

The first wall you hit when you start reading papers usually isn't the theory. It's shape.

The page says $\sum_j A_{ij} B_{jk}$. The reference implementation is full of x.transpose(1, 2).reshape(b, -1, d). You run it and get shapes (32,128,64) and (32,64,128) not aligned. You feel like you understood the math, but nothing connects it to the code. This is completely normal.

The bridge is einsum notation. Once you're used to it, a line like torch.einsum('bhqd,bhkd->bhqk', q, k) is easier to read than the paper's subscripts. Put the other way round: the moment einsum becomes readable, you can see that the $\sum$ in the paper and the transpose in the code are two spellings of one idea. This article gets you there from zero.

The metaphor: a warehouse where every axis has a label

"Tensor" means something precise in physics and differential geometry. In deep learning practice it means something much plainer: a box of numbers arranged in several dimensions. One number is a scalar, a row is a vector, a grid is a matrix, and past that we ran out of names — so everything else is just called a tensor.

What actually matters isn't the box, it's that each axis carries a meaning. A tensor of shape (32, 8, 128, 64) is not an anonymous pile of numbers:

• axis 0, 32 — batch (how many sequences we process at once)
• axis 1, 8 — head (how many viewpoints the attention layer has)
• axis 2, 128 — position (which token in the sequence)
• axis 3, 64 — feature (how many numbers represent one token)

As a warehouse: 32 buildings, 8 floors each, 128 drawers per floor, 64 compartments per drawer. Swap "building" and "floor" and the total number of slots is unchanged, so your program keeps running without complaint. That is exactly why shape bugs are silent.

So the first habit to build is: never read a shape as a list of numbers. When you see (32, 8, 128, 64), mentally relabel it (batch, head, position, feature). With that one habit, almost every operation below becomes "moving the labels around."

The intuition: there are only three things going on

Tensor code looks intricate, but the vocabulary is tiny.

1. Collapse an axis (contraction). — Sum along an axis and delete it. Sums, means, dot products and matrix products are all this. Collapse the last axis of (32, 128, 64) and you get (32, 128).

2. Keep or add an axis. — Don't delete it. Outer products, and carrying a batch dimension along untouched, are this.

3. Line axes up (broadcasting). — Stretch one operand's axes so two differently shaped tensors can be added or multiplied.

Beyond those there are only operations that move nothing and only reorder — transpose and reshape. Which means: any expression you meet in a paper can be read out in this vocabulary. einsum is a notation that packs all three into one string.

Lining up: broadcasting in three rules

Let's dispose of broadcasting first. Three lines:

1. Align the two shapes from the right (the shorter one is treated as having length-1 axes on the left).
2. At each position, the lengths must be equal, or one of them must be 1.
3. A length-1 axis behaves as if copied to match the other side (nothing is really copied — the same value is re-read).

An example:

x = np.zeros((32, 128, 64))   # (batch, position, feature)
b = np.zeros((64,))           # one bias per feature
x + b                         # fine → (32, 128, 64)

Aligned from the right, b is treated as (1, 1, 64), so the same 64 values are added at every batch-and-position cell. That is precisely the intent of "one bias per feature."

The dangerous case is a combination that satisfies the rules but not your intent:

a = np.zeros((1000,))      # (N,)
c = np.zeros((1000, 1))    # (N, 1)
a - c                       # legal → (1000, 1000)

You meant to subtract; you built a 1000×1000 distance matrix. No error is raised. This "a length-1 axis quietly multiplies" behaviour is broadcasting's biggest trap. When memory usage jumps for no visible reason, look here first.

The mechanism: deleting summation signs one at a time

Start from the definition of a matrix product.

Read aloud: "the value at row $i$, column $k$ of the output is row $i$ of $A$ and column $k$ of $B$, multiplied elementwise along the shared index $j$ and summed." Here $i$ and $k$ survive into the output; $j$ is summed away.

What Einstein noticed is that the summed index is obvious from looking at it: $j$ appears twice on the right and not at all on the left. So there's no need to write $\sum_j$. Dropping it gives a line that means exactly the same as equation (1):

That is the Einstein summation convention, and it's where the function name einsum comes from. Turn that line straight into a string and you get:

C = np.einsum('ij,jk->ik', A, B)

Commas separate the inputs; the output goes to the right of ->. The letters are yours to choose.

There are only three rules to read and write it:

• A letter to the right of -> is an axis that survives (a free index).
• A letter in the inputs but not in the output is summed over and disappears (contraction).
• The same letter in more than one input pairs those axes up and multiplies along them.

So you write down which axes you want to keep, and everything else collapses automatically. It clicks once you see it as a notation for what remains, not for what you do.

Practice reading and writing

The same rules cover every operation you'll actually use. Try reading each one out loud.

The behaviour worth memorising is b. A letter present in both inputs and in the output means "run the same computation independently for each slice of that axis." Batch, head, timestep — every axis you want to carry along in parallel is written this way.

With that vocabulary, the attention score computation reads:

# q, k: (batch, head, position, dim)
scores = torch.einsum('bhqd,bhkd->bhqk', q, k)   # (batch, head, q-position, k-position)

Out loud: "carry batch and head along, collapse the feature axis d, keep query position q and key position k." This is exactly $QK^\top$ — and because the string states which axis the dot product runs over, it's harder to misread than a version with transpose(-2, -1) buried in it. What attention is actually doing is covered in Attention from scratch.

The smallest unit — the dot product — is quicker to feel than to describe:

Writing it in code

Here's the body of attention, built out of einsum:

import torch

def attention(q, k, v, mask=None):
    # q, k, v: (b, h, n, d)
    d = q.shape[-1]
    scores = torch.einsum('bhqd,bhkd->bhqk', q, k) / d**0.5
    if mask is not None:
        scores = scores.masked_fill(mask, float('-inf'))  # mask: (b, 1, q, k), etc.
    w = scores.softmax(dim=-1)
    return torch.einsum('bhqk,bhkd->bhqd', w, v)          # weighted average

The two einsums are a matched pair. The first collapses the feature axis d to produce key positions k; the second collapses key positions k to bring the feature axis d back. The out-and-back structure is visible right there in the strings.

When you want to restructure axes rather than contract them, einops rearrange reads better than einsum:

from einops import rearrange
# split (b, n, h*d) into (b, h, n, d)
q = rearrange(q, 'b n (h d) -> b h n d', h=8)

The parentheses in (h d) say "this axis is h and d fused together." It does the same work as reshape(b, n, 8, -1).transpose(1, 2), but what was split into what is written down, so a reviewer can catch the mistake.

Contraction order changes the cost by orders of magnitude

einsum takes three or more inputs too:

np.einsum('ij,jk,kl->il', A, B, C)

Mathematically $(AB)C = A(BC)$, so the answer is identical. The cost is not. With $A$ of shape $(100, 1000)$, $B$ of $(1000, 1000)$ and $C$ of $(1000, 5)$:

• $(AB)C$: $100{\times}1000{\times}1000 + 100{\times}1000{\times}5 \approx 100.5$ million multiply-adds
• $A(BC)$: $1000{\times}1000{\times}5 + 100{\times}1000{\times}5 \approx 5.5$ million

More than an 18× difference, from nothing but the order of multiplication.

NumPy's np.einsum does not optimise the order by default. If you're contracting three or more operands at once, pass optimize=True, or inspect the chosen order with np.einsum_path:

print(np.einsum_path('ij,jk,kl->il', A, B, C, optimize='optimal')[1])

PyTorch's torch.einsum will use a path search when opt_einsum is installed, and you can control that through torch.backends.opt_einsum. Note that optimisation only has anything to do with three or more operands — with two, there is no order to choose. Where the cost of a single matrix product comes from is covered in The cost of matrix multiplication.

How this shows up on the job

Who uses it, and when. Research engineers turning papers into implementations; inference engineers reading somebody else's model code to speed it up; MLOps people chasing down why training won't converge. What they share is the need to read someone else's tensor manipulation correctly, fast. einsum and einops have become the shared language for that.

What you'll actually touch. torch.einsum / np.einsum / np.einsum_path / optimize=True; einops's rearrange, reduce, repeat; and while debugging, tensor.shape, tensor.stride(), tensor.is_contiguous().

The traps that cause incidents.

1. Using reshape to reorder axes. Going from (b, n, h, d) to (b, h, n, d) with reshape raises no error but scrambles values into the wrong places. Reordering is transpose / permute; reshape is only for fusing and splitting. It surfaces only as "training just sort of doesn't converge," which makes it extremely slow to find.

2. view blowing up. After a transpose the tensor is no longer contiguous in memory, so view throws. reshape succeeds — by silently making a copy.

3. Mask shapes. An attention mask of (b, 1, 1, k) (padding) means something different from (b, 1, q, k) (causal). Broadcasting accepts both, so the wrong one still runs — while letting the model see future tokens, which is the worst possible way to be broken.

4. Materialised intermediates. 'bhqd,bhkd->bhqk' allocates a $b{\times}h{\times}q{\times}k$ array in memory. Batch 8, 32 heads, sequence length 4096, fp16: $8{\times}32{\times}4096{\times}4096{\times}2$ bytes ≈ 8 GiB. The intermediate is larger than the output, and this is the classic cause of an OOM. Methods like FlashAttention are, in this framing, ways to avoid building that table at all.

5. Reusing index letters. Write both q and k in 'bhqd,bhkd->bhqk' as n and you have declared that queries and keys are the same axis. Give every distinct meaning its own letter.

What gets asked in interviews. "What does torch.einsum('bhqd,bhkd->bhqk', q, k) do?" is a standard question. (1) Point at the letter that collapses (d), (2) point at the ones that survive (b, h, q, k), (3) state it plainly — "batch and head ride along, the dot product runs over the feature axis, and the result is a position-by-position table." Say that much and the natural follow-up, "how many bytes is that table?", is just the multiplication above.

Summary

• A tensor is a box whose axes carry meaning. Read a shape as labels, not numbers.
• There are only three operations: collapse (contract), keep, and line up (broadcast).
• einsum describes what remains, not what you do. Any letter absent from the output disappears.
• With three or more operands, contraction order changes the cost by orders of magnitude. Don't forget optimize=True.
• Broadcasting and reshape go wrong without raising errors. That's what makes them dangerous.

If you want to go back to what vectors and matrices are doing in the first place, re-reading Linear algebra for AI makes einsum strings suddenly look like ordinary sentences.