Maitreya Suin
July 16, 2026
Start with a loss function that is wrong in a way you will not notice. Here is a batch of five predictions and their five targets, and the mean-squared error between them:
loss: 2.2934672832489014A number came out. It’s positive. It’s differentiable — call .backward() and gradients flow. Drop it into a training loop and the loss will go down, epoch after epoch, and you will believe everything is fine.
It is not fine. pred has shape (5,) and target has shape (5, 1), and when you subtract them, PyTorch does not complain about the mismatch — it broadcasts, producing a (5, 5) matrix in which every prediction is compared against every target. The loss you’re minimizing isn’t “how far is each prediction from its own target.” It’s “how far is each prediction from every target, on average,” and by the end of this post we’ll derive the exact consequence: this objective quietly pulls all five predictions toward the mean of the targets, and your model learns to predict a constant.
No exception was raised. No warning survived to your logs. The shapes were “compatible,” so PyTorch did what you asked — it just wasn’t what you meant.
This is the double edge of broadcasting, the feature that lets you add a bias vector to a matrix without a loop and also lets a squeezed tensor silently turn your regression into nonsense. This post is about seeing it clearly: the rule that governs when two shapes combine, a small grammar that most tensor computations fit into, the precise objective-and-gradient you get when the rule fires on axes that don’t line up semantically, and finally einsum — a notation where the axis bookkeeping is written down where a reviewer can check it.
TL;DR — Broadcasting aligns shapes from the right and logically stretches size-1 axes; that stretch is often just a stride-0 view, so it allocates nothing, but the operation consuming it still computes and stores a full-sized result. The danger is that “shapes compatible” and “shapes mean the same thing” are different tests, and only the first is automatic — so
(5,) − (5,1)gives a(5,5)outer difference with a perfectly valid gradient. Reductions are the complementary move (collapse an axis),keepdim=Trueleaves the size-1 stub that lets the result broadcast back, andeinsummakes the whole align-combine-reduce pattern explicit enough to review.
Before the rule, the reason the rule matters. A surprisingly large fraction of tensor code — not all of it, but enough that it’s worth having a mental template — decomposes into four moves:
A dot product is combine-then-reduce (a * b, then .sum()). Attention scores pair every query position with every key position and contract the feature axis within each head. A covariance matrix is center-combine-reduce. Even a convolution, if you squint past its weight-sharing, has this shape at its core. It’s not a universal law — plenty of real computation involves logarithms, sorting, indexing, normalization that don’t fit — but when you’re staring at an unfamiliar tensor expression, “which axes align, which broadcast, which get reduced” is the question that unlocks it. Broadcasting governs move 2; reductions govern move 4; and einsum is a notation for expressing 1–4 in a single string.
Two tensors are broadcast-compatible if, when you write their shapes right-aligned and pad the shorter one with leading 1s, every pair of dimensions is either equal or has a 1 in at least one slot. The output takes the larger size in each position.
a: (4, 1, 16)
b: (8, 16) -> padded to (1, 8, 16)
----------------------
a+b: (4, 8, 16) each axis: max of the twoFormally, after left-padding to equal rank, shapes are compatible when for every axis \(i\), \(a_i = b_i\) or \(a_i = 1\) or \(b_i = 1\); and the result size is \(c_i = \max(a_i, b_i)\).
The size-1 axes get stretched, and — connecting to the strides post — “stretched” usually means an internal stride-0 view: the same element read repeatedly, no memory allocated for the expansion. That last qualifier is the whole trap in miniature: the expansion is free, but the operation consuming it still evaluates every logical output position and allocates the result at the full broadcast shape.
The friendly warm-up is “compute a dot product without torch.dot,” and it’s friendly because a dot product is just move 3 followed by move 4: multiply elementwise, then sum the products away. For two 1-D vectors, every reasonable way of writing it agrees:
a, b = torch.randn(16), torch.randn(16)
ref = torch.dot(a, b)
for name, val in {
"(a * b).sum()": (a * b).sum(),
"(a * b).sum(dim=-1)": (a * b).sum(dim=-1),
"a @ b": a @ b,
"torch.inner(a, b)": torch.inner(a, b),
"einsum('i,i->', a, b)": torch.einsum('i,i->', a, b),
}.items():
print(f"{name:24s} {val.item():+.5f} matches torch.dot: {torch.allclose(val, ref)}")(a * b).sum() +4.73300 matches torch.dot: True
(a * b).sum(dim=-1) +4.73300 matches torch.dot: True
a @ b +4.73300 matches torch.dot: True
torch.inner(a, b) +4.73300 matches torch.dot: True
einsum('i,i->', a, b) +4.73300 matches torch.dot: TrueAll five match for these one-dimensional inputs — and that qualifier is load-bearing, because the moment you go to a batch of vectors and want one scalar per row, these operations stop being interchangeable. This is the genuinely useful thing to understand, and it trips up people constantly:
B, D = 4, 16
a, b = torch.randn(B, D), torch.randn(B, D)
print("row-wise dot, one scalar per row:")
print(" (a * b).sum(-1) ->", tuple((a * b).sum(-1).shape))
print(" torch.linalg.vecdot ->", tuple(torch.linalg.vecdot(a, b).shape))
print(" einsum('bd,bd->b') ->", tuple(torch.einsum('bd,bd->b', a, b).shape))
print("\nsame-named ops that do something DIFFERENT:")
print(" torch.inner(a, b) ->", tuple(torch.inner(a, b).shape), " <- every row vs every row")
print(" a @ b.T ->", tuple((a @ b.T).shape), " <- same: a full (B,B) matrix")row-wise dot, one scalar per row:
(a * b).sum(-1) -> (4,)
torch.linalg.vecdot -> (4,)
einsum('bd,bd->b') -> (4,)
same-named ops that do something DIFFERENT:
torch.inner(a, b) -> (4, 4) <- every row vs every row
a @ b.T -> (4, 4) <- same: a full (B,B) matrixHere’s the table worth remembering, because “dot product” names at least three different higher-rank operations:
| operation | (D,),(D,) |
(B,D),(B,D) |
|---|---|---|
(a*b).sum(-1) |
scalar | (B,) — one dot per row |
torch.linalg.vecdot(a, b) |
scalar | (B,) — one dot per row |
torch.inner(a, b) |
scalar | (B, B) — all pairs |
a @ b |
scalar | invalid (shapes don’t chain) |
a @ b.T |
— | (B, B) — all pairs |
einsum(...) |
whatever the equation says | whatever the equation says |
torch.linalg.vecdot is the one purpose-built for “one dot product per shared leading coordinate” — it contracts the last axis and broadcasts the rest, which is exactly the batched-dot semantics you usually want. torch.inner, despite the name, contracts only the last dimension and keeps both sets of leading dimensions, so two (B, D) inputs give you the (B, B) matrix of all cross-products — occasionally what you want, disastrously not what you want if you were expecting (B,). (One more wrinkle for completeness: on complex tensors vecdot conjugates its first argument, following the physics convention, while a bare (a*b).sum() does not — if you ever do dot products on complex data, that difference is real.)
The habit that helps, stated carefully: when your tensor’s contract puts features on the last axis, dim=-1 is more robust than a hard-coded positive index, because it survives an extra batch dimension appearing upstream (which it always eventually does). But dim=-1 is not a substitute for knowing which axis holds your features — it’s a good default given that knowledge, not a replacement for it.
keepdim existsReductions are move 4: they collapse one or more axes into a summary. The subtlety worth internalizing is keepdim, and its most important use here is leaving singleton axes that can broadcast back over the input.
x = torch.randn(4, 8)
print("x.sum(-1).shape ", tuple(x.sum(-1).shape))
print("x.sum(-1, keepdim=True).shape", tuple(x.sum(-1, keepdim=True).shape))
centered = x - x.mean(-1, keepdim=True) # subtract each ROW's own mean
print("\nevery row-mean now ~0?")
torch.testing.assert_close(centered.mean(-1), torch.zeros(x.shape[0]), atol=1e-6, rtol=0)
print(" yes (assert_close passed)")x.sum(-1).shape (4,)
x.sum(-1, keepdim=True).shape (4, 1)
every row-mean now ~0?
yes (assert_close passed)Without keepdim, x.mean(-1) has shape (4,), which does not line up against (4, 8) under right-alignment — the 4 would try to align with the 8 and fail. With keepdim, you get (4, 1), which broadcasts cleanly over the 8 columns. Many normalization formulas use a reduce-with-keepdim, transform, and broadcast-back pattern, although the chosen axes, statistics, state, and affine parameters differ.
Two reduction details that matter more than they first appear:
You can reduce several axes at once. Pass a tuple, and keepdim applies to all of them:
mean over the last two axes: (2, 3, 1, 1)mean over the last two axes: (2, 3, 1, 1)That’s the spatial-average-pool at the end of a vision model, in one call.
sum and mean differ by a factor of N in the gradient, and it’s not cosmetic. Summing \(N\) per-element losses gives gradients \(N\times\) larger than averaging them:
d(mean)/dx: [0.25, 0.25, 0.25, 0.25]
d(sum)/dx : [1.0, 1.0, 1.0, 1.0] <- N times largerd(mean)/dx: [0.25, 0.25, 0.25, 0.25]
d(sum)/dx : [1.0, 1.0, 1.0, 1.0] <- N times largerThis is why a loss reduction of sum vs mean silently changes your effective learning rate, and why it matters when averaging gradients across a distributed batch — get the reduction wrong and every device is effectively training at the wrong step size.
Everything so far shares a weakness: right-alignment matches sizes, and has no idea what an axis means. When two different axes happen to have the same size, broadcasting will cheerfully align the wrong ones and tell you nothing. The defense is a habit I’d genuinely recommend adopting for any non-trivial expression — write down a tiny ledger of what each axis is before you write the operation:
| tensor | shape | what each axis means |
|---|---|---|
q |
(B, H, Q, D) |
batch, head, query position, feature |
k |
(B, H, K, D) |
batch, head, key position, feature |
scale |
(1, H, 1, 1) |
broadcast over batch/query/key, one value per head |
| scores | (B, H, Q, K) |
batch, head, query, key |
With the ledger written, the operation reads off it directly, and — a small but real trick — giving every semantic axis a distinct size turns shape-checking into a free correctness test, because a mistake that swaps two axes can no longer hide behind them being equal:
B, H, Q, K, D = 2, 3, 5, 7, 11 # all different on purpose
q = torch.randn(B, H, Q, D)
k = torch.randn(B, H, K, D)
scale = torch.rand(1, H, 1, 1) # per-head scale, written to broadcast explicitly
scores = torch.einsum('bhqd,bhkd->bhqk', q, k) * scale
assert scores.shape == (B, H, Q, K)
print("scores:", tuple(scores.shape), "— and if I'd swapped Q/K, this assert would have caught it")scores: (2, 3, 5, 7) — and if I'd swapped Q/K, this assert would have caught it(That per-head scale is illustrative — standard attention divides by a single scalar \(\sqrt{D}\); per-head learned scales are a real but non-default variant. The point here is the shape discipline, not the specific scaling.)
Now we have the tools to finish the opening story properly — not just “it’s an outer product” but what objective and what gradient the bug creates. First, see it happen on numbers you can check in your head:
pred = torch.tensor([0., 10.]) # two predictions
target = torch.tensor([[0.], [10.]]) # two targets, shape (2, 1)
print("pred - target has shape:", tuple((pred - target).shape))
print("the (2,2) squared-error matrix:")
print(((pred - target) ** 2))
print("\nbuggy mean :", ((pred - target) ** 2).mean().item(), " <- pred EXACTLY matches its target, yet loss is 50")
print("correct mean:", ((pred - target.squeeze(-1)) ** 2).mean().item())pred - target has shape: (2, 2)
the (2,2) squared-error matrix:
tensor([[ 0., 100.],
[100., 0.]])
buggy mean : 50.0 <- pred EXACTLY matches its target, yet loss is 50
correct mean: 0.0Prediction 0 is target 0. Prediction 1 is target 1. The correct loss is zero — the model is perfect. But the broadcast version builds the full 2×2 matrix of every-pred-vs-every-target, and the off-diagonal (0−10)² = 100 terms drag the mean up to 50. The optimizer, handed this, will move the already-perfect predictions to reduce those off-diagonal terms, actively breaking a correct model.
We can write down exactly what it’s minimizing. The intended objective is
\[L_{\text{correct}} = \frac{1}{N}\sum_{j} (p_j - t_j)^2,\]
but broadcasting \((N,)\) against \((N,1)\) produces the double sum
\[L_{\text{bug}} = \frac{1}{N^2}\sum_{i}\sum_{j} (p_j - t_i)^2.\]
Differentiate that with respect to a single prediction \(p_j\) and the cross-terms collapse into something clean:
\[\frac{\partial L_{\text{bug}}}{\partial p_j} = \frac{2}{N}\left(p_j - \bar{t}\right), \qquad \bar{t} = \frac{1}{N}\sum_i t_i.\]
Every prediction is being pulled toward \(\bar t\), the mean of all targets — not toward its own target. That’s the whole pathology in one line: the buggy loss trains your model to ignore its input and output the average label (for a full batch). And because it’s a real objective with a real gradient, the loss curve slopes obligingly downward the entire time. Let’s confirm the gradient formula against autograd rather than trusting the algebra:
N = 5
pred = torch.randn(N, requires_grad=True)
target = torch.randn(N, 1)
((pred - target) ** 2).mean().backward()
from_formula = (2 / N) * (pred.detach() - target.mean())
torch.testing.assert_close(pred.grad, from_formula)
print("autograd gradient :", pred.grad.numpy().round(4))
print("(2/N)(p - mean(t)):", from_formula.numpy().round(4))
print("mean of targets :", target.mean().item(), "<- what every prediction is pulled toward")autograd gradient : [ 0.2784 -0.7326 0.3646 -0.5366 -0.2308]
(2/N)(p - mean(t)): [ 0.2784 -0.7326 0.3646 -0.5366 -0.2308]
mean of targets : 0.29763174057006836 <- what every prediction is pulled towardThe formula matches autograd exactly. This is the payoff of taking the bug seriously instead of just labeling it: a falling training curve does not validate the shape contract. The loss went down; the model got worse.
The fix is not to reshape defensively at the loss — it’s to make the shape contract explicit and fail loudly when it’s violated:
import torch.nn.functional as F
def mse_strict(pred, target):
"""Mean squared error that refuses to broadcast a shape mismatch into a silent bug."""
if pred.shape != target.shape:
raise ValueError(
f"pointwise MSE needs identical shapes; got pred={tuple(pred.shape)}, "
f"target={tuple(target.shape)}. Fix the shapes upstream, don't squeeze here."
)
return F.mse_loss(pred, target)Two deliberate choices in that guard. It raises a ValueError, not a Python assert — assertions vanish under python -O, so they’re fine for tests and wrong for a runtime contract you actually depend on. And it refuses rather than “helpfully” squeezing, because the right place to fix a shape mismatch is where the wrong shape was created (the dataframe column that came in as (N, 1)), not one line before the loss. If you do need to drop that trailing axis, use target.squeeze(-1) and never a bare target.squeeze() — the unqualified version removes every size-1 axis, so a (1, 1) target collapses to a scalar and you’ve traded one shape bug for another.
A corner of the broadcasting contract that foundations material usually skips, and shouldn’t: broadcasting an out-of-place operation computes a new shape, but an in-place operation can only broadcast the other operand into its own existing shape — it cannot grow the tensor it writes into.
out-of-place (x + y).shape: (3, 3, 7)--------------------------------------------------------------------------- RuntimeError Traceback (most recent call last) Cell In[14], line 5 2 y = torch.ones(3, 1, 7) 4 print("out-of-place (x + y).shape:", tuple((x + y).shape)) # (3, 3, 7), fine ----> 5 x.add_(y) # in-place: x can't grow to (3,3,7) RuntimeError: output with shape [1, 3, 1] doesn't match the broadcast shape [3, 3, 7]
The rule to carry: a + b produces the broadcast of both shapes; a += b requires that broadcast to fit inside a’s existing shape. The in-place form will happily broadcast b up to a, but it will not let a itself expand — because it has nowhere to put the extra elements. This bites when you write buffer += update expecting a convenient broadcast and get an error, and the fix is usually to preallocate buffer at the full shape.
einsum gives a compact notation for axis alignment, multiplication, and summation-based contraction. Four rules cover almost everything:
Rule 1 — one label per input axis. 'bhqd,bhkd->bhqk': the first operand has axes b,h,q,d; the second has b,h,k,d.
Rule 2 — a shared label aligns those axes. The two bs are the same batch coordinate; the two hs the same head; the two ds the same feature axis, to be contracted.
Rule 3 — a label absent from the output is summed away. d appears in both inputs but not the output, so it’s the contraction axis: multiply along it and sum.
Rule 4 — a label repeated within one operand takes a diagonal, not a contraction. 'ii->i' is the diagonal of a matrix; 'ii->' is the trace. This is the one clause a one-sentence summary of einsum always gets wrong, and it’s why “repeated label = summed” is not quite the rule.
q = torch.randn(2, 3, 5, 11) # (B, H, Q, D) — distinct sizes, so shape checks are meaningful
k = torch.randn(2, 3, 7, 11) # (B, H, K, D)
M = torch.randn(4, 4)
# attention scores, as itself:
scores = torch.einsum('bhqd,bhkd->bhqk', q, k)
print("scores 'bhqd,bhkd->bhqk':", tuple(scores.shape))
torch.testing.assert_close(scores, q @ k.transpose(-2, -1)) # exact, useful failure if wrong
print(" matches q @ k.transpose(-2,-1) within floating-point tolerance")
print("diagonal 'ii->i':", torch.einsum('ii->i', M).shape, "| trace 'ii->':",
torch.allclose(torch.einsum('ii->', M), torch.trace(M)))scores 'bhqd,bhkd->bhqk': (2, 3, 5, 7)
matches q @ k.transpose(-2,-1) within floating-point tolerance
diagonal 'ii->i': torch.Size([4]) | trace 'ii->': TrueRead 'bhqd,bhkd->bhqk' aloud: for each batch and head, contract the feature axis d, keeping query q against key k. That is the definition of an attention score matrix, transcribed. Compare it to q @ k.transpose(-2, -1) — the same computation, but expressed as “perform a transpose you now have to verify in your head, and trust the batch dims chain correctly.” The einsum names its axes; the matmul makes you track them.
Two things that make einsum more reusable, and one warning:
Prefer explicit output notation. Write 'ij,jk->ik', not the implicit 'ij,jk'. Spelling out the output removes ambiguity and documents intent.
Ellipsis handles leading axes you don’t want to name. '...qd,...kd->...qk' computes attention scores over whatever batch-like dimensions precede q,d — batch, head, ensemble, beam, any combination:
ellipsis '...qd,...kd->...qk': (2, 3, 5, 7) (same result, fewer assumptions)The warning, which is the whole reason einsum is a tool and not a guarantee: a valid equation can be the wrong equation. einsum makes your assumed contract visible; it does not check that the contract is right.
wrong 'bhqd,bhkd->bhkq': (2, 3, 7, 5) — runs fine, it's the transpose of what you wanted
silently equals scores.transpose(-1,-2)? TrueSwapping the last two output labels transposes your scores, runs without complaint, and — if Q and K had been the same size — would have produced the right shape with the wrong contents, the exact same class of bug as the opening loss. Distinct axis sizes are what surfaced it here; that’s the cheap-adversarial-test point from the ledger section, doing real work.
The whole point of the align-combine-reduce grammar is that seemingly different computations turn out to be the same shape of thing. A sample covariance matrix is center-then-contract — the same moves as attention scores, on different axes:
def sample_covariance(x):
"""Sample covariance of observations x with shape (N, D): N observations, D features.
Rows are observations, columns are features; uses Bessel's (N-1) correction."""
if x.ndim != 2:
raise ValueError(f"expected (N, D); got {tuple(x.shape)}")
n = x.shape[0]
if n < 2:
raise ValueError("covariance needs at least 2 observations")
centered = x - x.mean(dim=0, keepdim=True) # move 1+2: center via keepdim-broadcast
return (centered.transpose(0, 1) @ centered) / (n - 1) # move 3+4: contract the observation axis
obs = torch.randn(100, 5)
cov = sample_covariance(obs)
print("covariance shape:", tuple(cov.shape), "| symmetric?",
torch.allclose(cov, cov.T, atol=1e-6))covariance shape: (5, 5) | symmetric? TrueI wrote that as a full function with its assumptions checked rather than the one-liner x.T @ x / (n-1), because the one-liner hides everything that can go wrong: it assumes rows-are-observations, N > 1, real-valued data, and Bessel’s correction, and it references an undefined n. Making the contract explicit is the same discipline as mse_strict — the correctness-first version states what it needs and refuses when it doesn’t get it.
If you can see that this covariance, the attention scores above, and the batched dot product from earlier are all align two things on a shared axis, combine, contract that axis away, then the grammar has done its job — you’re no longer memorizing separate operations, you’re recognizing one pattern in different clothes.
Broadcasting is the feature that makes tensor code concise and the feature that makes it silently wrong, and both come from the same mechanism: PyTorch aligns shapes from the right and computes whatever is compatible, whether or not the aligned axes mean the same thing.
(N,) − (N,1) pulls every prediction toward the mean of the targets.einsum is worth using because the axis contract becomes a string a reviewer can read — but a valid equation can still be the wrong equation, so it’s a tool for visibility, not a proof of correctness.If one sentence survives: a falling loss curve is not evidence that your shapes mean what you think they do — check the axis contract, because the machine only checks the sizes.
We’ve now seen the “same shape, wrong meaning” trap in two forms — the packed-axis bug of the strides post and the broadcast-objective bug here. The next post leaves correctness for performance: how the GPU actually executes these operations, where “broadcasting allocates the result” and “contiguity has a cost” stop being abstract and get concrete, measured numbers — and why a masking operation that looks free can quietly force a synchronization that stalls the whole pipeline.
torch.linalg.vecdot — batched dot products, and the complex-conjugation convention that separates it from a bare (a*b).sum().torch.einsum (ellipsis, repeated-label semantics) and Olexa Bilaniuk’s einsum tutorial for intuition.torch.sum — the dtype= accumulation argument and integer-promotion behavior.opt_einsum — contraction-order optimization, relevant at three-plus operands.