Shapes, batches, and the broadcasting rules

The batch dimension is a convenience, not a concept

Real code never multiplies one vector at a time. It stacks 32 or 256 examples into a batch and processes them together, because a GPU is a machine for doing the same arithmetic to many things at once. So every tensor grows a leading dimension, and the shapes you read look like (32, 128, 768): 32 examples, 128 token positions each, 768 features each.

The rule that keeps this simple is that matrix multiplication in NumPy and PyTorch acts on the last two dimensions and treats everything in front as batch. So (32, 128, 768) @ (768, 3072) gives (32, 128, 3072): the same 768 × 3072 weight matrix applied to each of the 32 × 128 = 4,096 individual vectors, with no loop written anywhere.

Broadcasting, in three rules

Broadcasting is what lets you write x + b where x is (32, 128, 3072) and b is just (3072,). The rules, applied right to left:

1. Line the shapes up from the right, padding the shorter one with 1s on the left.
2. Two dimensions are compatible if they are equal, or if one of them is 1.
3. A dimension of size 1 is stretched — conceptually repeated — to match the other.

Worked:

x   (32, 128, 3072)
b               (3072,)   ->  (1, 1, 3072) after padding
                              every axis compatible
result (32, 128, 3072)

The bias vector is added to every position of every example. No memory is actually copied; the library reads the same 3,072 numbers repeatedly. That is why broadcasting is fast as well as short.

The bug that costs a week

Now the failure. Broadcasting is too accommodating: it will happily produce a valid result from shapes you did not intend to combine.

import numpy as np
predictions = np.array([1.0, 2.0, 3.0, 4.0])       # shape (4,)
targets     = np.array([[1.1], [2.1], [2.9], [4.2]])  # shape (4, 1)
error = predictions - targets
print(error.shape)     # (4, 4)  -- not (4,)
print(error.mean())    # a number, and completely meaningless

You wanted four errors. You got sixteen, because (4,) padded to (1, 4) and (4, 1) broadcast against each other into a full 4×4 grid of every prediction against every target. The mean is a real number. The loss decreases during training. The model learns something, slowly and wrongly, and nothing raises an exception.

This is the single most common silent bug in hand-written training code, and it comes from a column vector where a flat vector was expected — most often straight out of a pandas column or a reshape(-1, 1) somebody added to satisfy a different function.

The defence is one line:

assert predictions.shape == targets.shape, (predictions.shape, targets.shape)

Put it before every loss computation. It costs nothing and it catches this class of bug entirely.

The shape operations, and what each really does

Four operations show up constantly, and confusing them causes the rest of the shape bugs.

• reshape reinterprets the same numbers in the same memory order under a new shape. (2, 6) to (3, 4) is legal because both hold 12 numbers. It never moves data between positions in the flat ordering.
• transpose genuinely reorders axes. (32, 128, 768) transposed to (32, 768, 128) puts different numbers in different places. Reshape and transpose are not interchangeable, and swapping them is the classic bug when splitting attention heads.
• squeeze removes dimensions of size 1; unsqueeze (or None indexing) inserts one. These exist mainly to make broadcasting do what you meant.
• view versus copy. Reshape usually returns a view sharing memory with the original, so writing to one changes the other. Transpose does too. This is fast and occasionally surprising.

Getting the head split right

The place this all comes together is multi-head attention. A tensor of shape (batch, seq, d_model) with d_model = 768 and 12 heads has to become (batch, heads, seq, d_head) with d_head = 64. The correct sequence is reshape then transpose:

x = x.reshape(batch, seq, 12, 64)      # split the feature axis
x = x.transpose(0, 2, 1, 3)            # (batch, heads, seq, d_head)

Reshaping straight to (batch, 12, seq, 64) is legal arithmetic and gives the wrong answer: it slices the sequence across heads instead of slicing the features. The model still trains. It just learns a worse function, and you have no error to chase.

The rule to keep

Every dimension of size 1 is an invitation for the library to stretch it. That is a feature when you meant it and a silent bug when you did not, so assert the shapes you expect rather than trusting that a result which computed is a result which is correct.