Shapes, broadcasting and the bugs they cause

Nearly every PyTorch bug is a shape bug, and the dangerous ones raise no error at all: subtract a tensor of shape [N] from one of shape [N, 1] and PyTorch quietly compares every prediction with every target, trains happily, and hands you a useless model. This section is about reading and predicting shapes, the broadcasting rules that decide when two tensors may meet, the difference between view, reshape and permute, and einsum once more than two dimensions are in play. The habit it is trying to install is writing down the shape of every tensor as you create it.

Say the shape out loud#

Most PyTorch bugs are shape bugs. The habit that prevents them: write the shape of every tensor as a comment when you create it, in named dimensions such as batch, time and channels.

import torch

x = torch.randn(32, 128, 256)      # [batch, time, d_model]
w = torch.randn(256, 64)           # [d_model, d_head]
h = x @ w                          # [batch, time, d_head] -> [32, 128, 64]

Broadcasting#

When two tensors of different shapes meet in an operation, PyTorch lines their shapes up from the right. Each pair of dimensions must be equal, or one of them must be 1, which is then stretched. Missing dimensions on the left count as 1.

[32, 128, 256]  +  [256]         ->  [32, 128, 256]   bias added to every position
[32, 1, 256]    +  [1, 128, 256]  ->  [32, 128, 256]
[32, 128]       +  [32]           ->  error: 128 and 32 do not match

Two rules produce all of it. Missing dimensions are added on the left as 1, and any dimension of size 1 is stretched to match the other. Nothing is copied while doing it, which is why broadcasting is fast and also why a [1000, 1] against a [1000] silently allocates a million elements without warning you that it has.

The bug that does not crash#

pred = model(x)                   # [N, 1]
y = targets                       # [N]
loss = ((pred - y) ** 2).mean()   # [N, 1] - [N] broadcasts to [N, N]. No error.

loss = ((pred.squeeze(-1) - y) ** 2).mean()   # what you meant

The first version compares every prediction with every target. It trains, the loss goes down a little, and the model is useless. Assert shapes at the boundaries of your functions to catch this.

The reason this survives review is that the number it prints is entirely plausible. Compare the two on random data:

torch.manual_seed(0)
pred = torch.randn(1000, 1)
y = torch.randn(1000)

print(tuple((pred - y).shape))               # (1000, 1000)  a million elements
print(((pred - y) ** 2).mean().item())               # 1.9996
print(((pred.squeeze(-1) - y) ** 2).mean().item())   # 2.0374

And it does train, which is the part that costs a day. Fitting a linear model to a problem it can solve exactly, with and without the squeeze:

for squeeze in (True, False):
    torch.manual_seed(0)
    X = torch.randn(500, 4); w = torch.randn(4, 1)
    Y = (X @ w).squeeze(-1) + 0.1 * torch.randn(500)
    m = nn.Linear(4, 1); opt = torch.optim.SGD(m.parameters(), lr=0.05)
    for _ in range(300):
        p = m(X)
        loss = ((p.squeeze(-1) - Y) ** 2).mean() if squeeze else ((p - Y) ** 2).mean()
        opt.zero_grad(); loss.backward(); opt.step()
    print(squeeze, round(((m(X).squeeze(-1) - Y) ** 2).mean().item(), 4))

# squeeze=True    true mse 0.0102      solved
# squeeze=False   true mse 4.0588      400 times worse, no error, loss fell the whole way

So the defence cannot be watching the loss. It has to be an assertion at the point the shapes meet, which costs one line and no measurable time.

def mse(pred, target):
    assert pred.shape == target.shape, f"{tuple(pred.shape)} vs {tuple(target.shape)}"
    return ((pred - target) ** 2).mean()

# Better still, make the intent explicit rather than relying on a squeeze:
loss = mse(pred.reshape(-1), target.reshape(-1))

Reshaping#

• view changes the shape without copying, and only works when the memory layout allows it.
• reshape does the same when it can and copies when it must.
• permute and transpose reorder dimensions by changing strides, which leaves the tensor non-contiguous. A view after them fails; call contiguous() first or use reshape.

The underlying idea is the stride: the number of elements to step to move one place along each dimension. A tensor is a flat block of memory plus a shape and a stride, so transpose does not move a single number, it swaps two entries in the stride tuple. That is why it is free, and why the result no longer matches the layout view assumes.

q = torch.randn(32, 128, 8 * 64)                  # [batch, time, heads * d_head]
q = q.view(32, 128, 8, 64).transpose(1, 2)          # [batch, heads, time, d_head]
out = q.transpose(1, 2).reshape(32, 128, 8 * 64)    # back: reshape, since transpose broke contiguity

t = torch.randn(32, 128, 512).view(32, 128, 8, 64).transpose(1, 2)
print(tuple(t.shape))        # (32, 8, 128, 64)
print(t.is_contiguous())     # False
print(t.stride())            # (65536, 64, 512, 1)  <- dim 1 steps 64, dim 2 steps 512

t.view(32, 128, 512)
# RuntimeError: view size is not compatible with input tensor's size and stride
#               (at least one dimension spans across two contiguous subspaces)

t.reshape(32, 8 * 128, 64)        # works: copies when it must
t.contiguous().view(32, 8, 8192)  # works: copy first, then a free view

So reshape is not simply the safer view. reshape returns a view when it can and a copy when it cannot, which means it never raises but its cost is unpredictable, and the returned tensor may or may not share memory with the original. view always shares and always raises if it cannot. Use view where you want to be told about an accidental copy in a hot loop, and reshape where you only want the result.

einsum#

For anything beyond a plain matrix multiply, einsum names every dimension and says which are summed. It is often the clearest way to write attention.

scores = torch.einsum("bhqd,bhkd->bhqk", q, k) / d_head ** 0.5
# d is summed away; b, h, q, k remain

The rule is short: name every dimension of every input, then name the ones you want out. Any letter that appears on the left and not on the right is summed over. Nothing else is going on, and it makes the alternative, a sequence of transposes and unsqueezes whose correctness you have to reconstruct each time you read it, unnecessary.

q = torch.randn(2, 4, 16, 32); k = torch.randn(2, 4, 16, 32)

a = torch.einsum("bhqd,bhkd->bhqk", q, k)
b = q @ k.transpose(-2, -1)
print((a - b).abs().max().item())   # 0.0   the same operation, written twice
print(tuple(a.shape))               # (2, 4, 16, 16)

einsum string            What it does
"ij,jk->ik"              plain matrix multiply
"bij,bjk->bik"           batched matrix multiply
"bhqd,bhkd->bhqk"        attention scores: sum over the head dimension d
"bhqk,bhkd->bhqd"        apply those scores to the values
"bi,bi->b"               row-wise dot product, without a sum and a multiply
"bij->bji"               transpose
"ii->i"                  the diagonal
"bchw,bcHW->bhwHW"       when you have gone too far and should draw it instead

dtypes#

• Class labels for cross_entropy must be integers (long), not floats.
• Default floats are 32-bit. Mixed precision uses 16-bit for speed; bfloat16 has the same range as 32-bit and rarely overflows, float16 can.
• Mixing dtypes in one operation raises an error or silently upcasts. Check with .dtype when something is odd.