Tensor shapes and broadcasting

This sheet is for anyone who has stared at RuntimeError: shape mismatch or, worse, at a model that trains fine and is silently wrong because a tensor broadcast where it should have errored. The rules are the same in NumPy, PyTorch and JAX. Blocks run on a CPU in under a second.

Broadcasting rules

Two shapes are compared from the right. Each pair of dimensions must be equal, or one of them must be 1 (or missing). The result takes the larger size in each position. There are no other rules.

Left	Right	Aligned	Result	Why
`(8, 1, 4)`	`(5, 1)`	`(8, 1, 4)` and `(1, 5, 1)`	`(8, 5, 4)`	Missing dims become 1, then 1 stretches
`(3, 4)`	`(4,)`	`(3, 4)` and `(1, 4)`	`(3, 4)`	A row vector adds to every row
`(3, 4)`	`(3,)`	`(3, 4)` and `(1, 3)`	error	4 and 3 differ and neither is 1. Use `(3, 1)`
`(3, 4)`	`(3, 1)`	same	`(3, 4)`	A column vector adds to every column
`(32, 10)`	`(32,)`	`(32, 10)` and `(1, 32)`	error unless the batch is 10	The classic per-sample weight bug. Use `w[:, None]`
`(2, 1, 5, 5)`	`(1, 4, 5, 5)`	same	`(2, 4, 5, 5)`	Mask over batch and head dims
`(N, C)`	`(C,)`	`(N, C)` and `(1, C)`	`(N, C)`	Bias or per-class scale

import numpy as np

a = np.ones((8, 1, 4))      # rank 3
b = np.ones((5, 1))         # rank 2, right-aligned to (1, 5, 1)
c = a + b                   # (8, 1, 4) with (1, 5, 1) -> (8, 5, 4)
print("a", a.shape, "+ b", b.shape, "->", c.shape)

x = np.arange(12.0).reshape(3, 4)
mean_per_row = x.mean(axis=1, keepdims=True)   # (3, 1), broadcasts back over columns
print("centred rows:", (x - mean_per_row).shape)

try:
    np.ones((3, 4)) + np.ones((4, 3))
except ValueError as err:
    print("error:", err)

a (8, 1, 4) + b (5, 1) -> (8, 5, 4)
centred rows: (3, 4)
error: operands could not be broadcast together with shapes (3,4) (4,3)

keepdims=True is the habit that avoids most broadcasting bugs: a reduction keeps its rank, so the result lines up with the input without a None or unsqueeze. The dangerous case is when the shapes happen to line up by accident, for example a (32,) per-sample weight against (32, 32) pairwise scores, which broadcasts across the wrong axis and trains.

Batch dimensions and the conventions you have to remember

Data	PyTorch layout	Note
Tabular	`(N, features)`
Sequence of tokens	`(N, T)` ids, `(N, T, D)` embeddings	`batch_first=True` in `nn.LSTM` and `nn.MultiheadAttention`, otherwise `(T, N, D)`
Images	`(N, C, H, W)`	PIL and most image libraries give `(H, W, C)`. TensorFlow defaults to `(N, H, W, C)`
Audio	`(N, channels, samples)` or `(N, T, mel_bins)` after a spectrogram
Attention scores	`(N, heads, T_q, T_k)`	Masks are `(T_q, T_k)` or `(N, 1, T_q, T_k)`
Logits for classification	`(N, classes)`	Labels `(N,)` of `int64`, not one-hot
Logits for token prediction	`(N, T, vocab)`	Flatten to `(NT, vocab)` and `(NT,)` for `cross_entropy`, or pass `(N, vocab, T)`

einsum patterns

einsum names each axis with a letter, repeats letters that should be summed over, and lists the output order. It replaces chains of transpose, unsqueeze and bmm with one line you can read back as a shape statement.

Pattern	Reads as	Equivalent
`"ij,jk->ik"`	matrix product	`a @ b`
`"bij,bjk->bik"`	batched matrix product	`torch.bmm(a, b)`
`"btd,de->bte"`	linear layer on every token	`x @ w`
`"bhqd,bhkd->bhqk"`	attention scores, all heads at once	`q @ k.transpose(-2, -1)`
`"bhqk,bhkd->bhqd"`	attention-weighted sum of values	`attn @ v`
`"bi,bj->bij"`	batched outer product	`a[:, :, None] * b[:, None, :]`
`"ii->"`	trace	`torch.trace(a)`
`"bij->bji"`	batched transpose	`a.transpose(1, 2)`
`"bd,bd->b"`	row-wise dot product	`(a * b).sum(-1)`

import torch

torch.manual_seed(0)
B, T, D, H = 2, 5, 8, 2          # batch, tokens, model dim, heads
x = torch.randn(B, T, D)
w = torch.randn(D, D)

# matrix multiply: (B,T,D) x (D,D) -> (B,T,D)
y = torch.einsum("btd,de->bte", x, w)

# split heads: (B,T,D) -> (B,H,T,D/H)
q = y.view(B, T, H, D // H).transpose(1, 2)
k, v = q.clone(), q.clone()

# attention scores per head: (B,H,T,d) x (B,H,T,d) -> (B,H,T,T)
scores = torch.einsum("bhqd,bhkd->bhqk", q, k) / (D // H) ** 0.5
attn = scores.softmax(dim=-1)

# weighted sum of values: (B,H,T,T) x (B,H,T,d) -> (B,H,T,d)
out = torch.einsum("bhqk,bhkd->bhqd", attn, v)
merged = out.transpose(1, 2).reshape(B, T, D)

# the same with matmul, to check
scores2 = (q @ k.transpose(-2, -1)) / (D // H) ** 0.5
print("scores", scores.shape, "merged", merged.shape, "einsum == matmul:", torch.allclose(scores, scores2))

# batched outer product and trace are one line each
outer = torch.einsum("bi,bj->bij", torch.randn(3, 4), torch.randn(3, 6))
trace = torch.einsum("ii->", torch.eye(4))
print("outer", outer.shape, "trace", trace.item())

scores torch.Size([2, 2, 5, 5]) merged torch.Size([2, 5, 8]) einsum == matmul: True
outer torch.Size([3, 4, 6]) trace 4.0

Note the .reshape after the final transpose: the heads were moved next to each other in logical order but not in memory. That is the next section.

Reshape versus permute

reshape and view keep the memory order and change how it is divided into dimensions. permute and transpose change which axis is which. They are not interchangeable, and confusing them produces tensors with the right shape and scrambled contents, which no error will catch.

import torch

t = torch.arange(6).reshape(2, 3)
print(t)
print("reshape(3, 2):\n", t.reshape(3, 2))     # reads elements in memory order, row 0 then row 1
print("permute(1, 0):\n", t.permute(1, 0))     # transposes: element [i, j] moves to [j, i]

p = t.permute(1, 0)
print("permuted is contiguous:", p.is_contiguous())
try:
    p.view(-1)
except RuntimeError as err:
    print("view fails:", str(err).split(".")[0])
print("contiguous().view works:", p.contiguous().view(-1).tolist())
print("reshape copies when needed:", p.reshape(-1).tolist())

reshape(3, 2) gives rows [0, 1], [2, 3], [4, 5]. permute(1, 0) gives rows [0, 3], [1, 4], [2, 5]. Same shape, different tensors. After a permute the tensor is non-contiguous, view refuses, and reshape quietly copies.

You want	Use	Not
Image `(H, W, C)` to `(C, H, W)`	`img.permute(2, 0, 1)`	`img.reshape(C, H, W)` (scrambles pixels)
Merge heads `(B, H, T, d)` to `(B, T, H*d)`	`x.transpose(1, 2).reshape(B, T, H * d)`	`x.reshape(B, T, H * d)` (mixes tokens and heads)
Flatten a batch of images for a linear layer	`x.flatten(1)`	`x.view(-1)` (loses the batch axis)
Add a batch axis	`x.unsqueeze(0)` or `x[None]`	`x.reshape(1, -1)` (also flattens)
Swap sequence and batch	`x.transpose(0, 1)`	`x.reshape(T, N, D)`

Shape bugs in attention and convolutions

import torch
from torch import nn

torch.manual_seed(0)

# Convolutions: PyTorch wants (N, C, H, W). An image loaded as (H, W, C) has to be permuted.
img_hwc = torch.rand(32, 32, 3)
x = img_hwc.permute(2, 0, 1).unsqueeze(0)       # (1, 3, 32, 32)
conv = nn.Conv2d(3, 8, kernel_size=3, stride=2, padding=1)
out = conv(x)
assert out.shape == torch.Size([1, 8, 16, 16]), out.shape
print("conv out", tuple(out.shape))             # floor((32 + 2 - 3) / 2) + 1 = 16

# Attention: a mask of the wrong rank broadcasts silently. Pin the shapes.
B, H, T = 2, 4, 6
scores = torch.randn(B, H, T, T)
causal = torch.tril(torch.ones(T, T, dtype=torch.bool))        # (T, T)
assert causal.shape == (T, T)
masked = scores.masked_fill(~causal, float("-inf"))            # broadcasts over B and H
probs = masked.softmax(dim=-1)
assert torch.allclose(probs.sum(-1), torch.ones(B, H, T))
print("row 0 attends to", int((probs[0, 0, 0] > 0).sum()), "token(s); row 5 attends to", int((probs[0, 0, 5] > 0).sum()))

# Loss functions have shape contracts too
logits = torch.randn(8, 10)            # (N, classes)
labels = torch.randint(0, 10, (8,))    # (N,) integer class ids, not one-hot
print("cross-entropy", round(nn.functional.cross_entropy(logits, labels).item(), 3))

conv out (1, 8, 16, 16)
row 0 attends to 1 token(s); row 5 attends to 6
cross-entropy 2.883

The convolution output size is floor((H + 2 * padding - dilation * (kernel - 1) - 1) / stride) + 1 per spatial axis. Write it in a comment next to the layer. The asserts cost nothing at runtime and turn a scrambled result into a traceback.

Bug	Symptom	Check
Padding mask shaped `(N, T)` applied where `(N, 1, 1, T)` is needed	Error, or masking the wrong axis	`mask[:, None, None, :]` and assert rank 4
Causal mask built with `triu` instead of `tril`	Model attends to the future, loss suspiciously low	`assert mask[0, -1] and not mask[-1, 0]` on the lower-left and upper-right corners
Softmax over `dim=-2`	Columns sum to one instead of rows	`assert torch.allclose(probs.sum(-1), 1)`
`nn.MultiheadAttention` fed `(N, T, D)` without `batch_first=True`	Batch and time swapped, trains badly	Pass `batch_first=True` and assert output shape equals input shape
Channels-last image into `Conv2d`	Error if C is not 3, or garbage if H happens to be 3	`assert x.shape[1] == in_channels`
Label shape `(N, 1)` with `cross_entropy`	Error, or with `BCEWithLogitsLoss` a `(N, 1)` vs `(N,)` broadcast to `(N, N)` warning	`labels.squeeze(-1)` and assert `labels.shape == logits.shape[:1]`
Batch of size 1 squeezed away by `.squeeze()`	Works until the last batch	`squeeze(dim)` with an explicit axis

torch.Size checks as assertions

tensor.shape is a torch.Size, a tuple subclass, so it compares to plain tuples. A one-line assert after each stage of a model is cheaper than any debugger. Keep them in the code: they document the contract and cost microseconds. For shapes that vary, assert the rank and the fixed dims, for example assert x.ndim == 3 and x.shape[-1] == d_model. For a longer-lived contract, name the dimensions in the function signature docstring, or use the jaxtyping annotations, which document shapes and can be checked at run time.