Broadcasting rules and the silent shape bug

How broadcasting lines up shapes

Broadcasting lets you combine tensors of different shapes without copying data. PyTorch uses the same rules as NumPy. Line the shapes up from the rightmost dimension. Two sizes are compatible if they are equal, or one of them is 1, or one tensor has no dimension there. Size-1 and missing dimensions are virtually stretched to match.

Examples: (32, 10) + (10,) gives (32, 10), which is how a bias vector is added to every row of a batch. (3, 1) + (1, 4) gives (3, 4), an outer-sum table. (4, 3) + (4,) fails, because the trailing sizes 3 and 4 differ and neither is 1.

The danger is that broadcasting succeeds when you did not intend it. A model that outputs shape (N, 1) compared with targets of shape (N,) broadcasts to (N, N): every prediction is compared with every target. No exception is raised; the loss is simply wrong and training quietly underperforms. nn.MSELoss prints a warning for this case, but many hand-written losses do not.

The silent (N, 1) vs (N,) bug

Shown with NumPy because the rules are identical and it runs anywhere. The wrong version raises no error; it just averages over a 3 x 3 grid of pairs.

import numpy as np  # PyTorch follows the same broadcasting rules

preds = np.array([[0.9], [0.2], [0.4]])   # shape (3, 1), like a model output
targets = np.array([1.0, 0.0, 0.0])       # shape (3,)

diff = preds - targets                    # silently broadcasts to (3, 3)!
print(diff.shape)
print(round(((preds.squeeze(1) - targets) ** 2).mean(), 4))  # correct
print(round((diff ** 2).mean(), 4))                          # wrong, no error

Output:

(3, 3)
0.07
0.3367

Assert shapes at boundaries

Before computing a loss, add assert preds.shape == targets.shape, (preds.shape, targets.shape). In PyTorch, fix the mismatch with preds.squeeze(1) or targets.unsqueeze(1). One assert line saves hours of wondering why accuracy is stuck.

त्वरित जाँच: What is the result shape of adding tensors with shapes `(4, 1, 3)` and `(5, 1)`?

• (4, 5, 3)
• It raises an error
• (4, 1, 3)
• (5, 4, 3)