Broadcasting in PyTorch With Clear Rules

Broadcasting is how PyTorch performs element-wise operations on tensors of different shapes without copying data. The smaller tensor is virtually stretched to match the larger one. It is powerful and compact, but it can also hide bugs, so you need to know the rules exactly.

The two rules

Line up the shapes by their trailing (rightmost) dimensions. Then, for each pair of dimensions:

1. If the sizes are equal, they are compatible.
2. If one of the sizes is 1, it is stretched to match the other.
3. Otherwise PyTorch raises an error. Missing leading dimensions count as size 1.

import torch

a = torch.ones(3, 4)
b = torch.tensor([1.0, 2.0, 3.0, 4.0])   # shape (4,)

print((a + b).shape)    # (3, 4): b is added to each row

Aligned on the right, (3, 4) and (4,) match in the last dimension, and the missing leading dimension of b is treated as 1 and stretched to 3. The result is (3, 4).

A column versus a row

col = torch.tensor([[10.0], [20.0], [30.0]])   # shape (3, 1)
row = torch.tensor([1.0, 2.0, 3.0, 4.0])        # shape (4,)

print((col + row).shape)   # torch.Size([3, 4])
print(col + row)
# tensor([[11., 12., 13., 14.],
#         [21., 22., 23., 24.],
#         [31., 32., 33., 34.]])

Here (3, 1) meets (1, 4) (after padding) and both are stretched to (3, 4). This outer-sum trick is useful for building grids and pairwise comparisons.

Practical use: normalising features

x = torch.rand(100, 5)                 # 100 samples, 5 features

mean = x.mean(dim=0)                   # shape (5,)
std = x.std(dim=0)                     # shape (5,)
x_norm = (x - mean) / std              # broadcasts over the 100 rows

print(x_norm.mean(dim=0))              # close to 0
print(x_norm.std(dim=0))               # close to 1

Per-feature statistics have shape (5,), which broadcasts across all rows with no loops. This is the same idea behind batch normalisation.

Adding dimensions with unsqueeze or None

When the shapes do not line up the way you want, add a size-1 dimension.

points = torch.rand(5, 2)              # 5 points in 2-D

diff = points.unsqueeze(1) - points.unsqueeze(0)   # (5,1,2) - (1,5,2)
print(diff.shape)                      # torch.Size([5, 5, 2])

dist = diff.pow(2).sum(dim=-1).sqrt()  # (5, 5) pairwise distances
print(dist.shape)
print(dist.diagonal())                 # all zeros

This computes every pairwise distance in three lines. The same result is available from torch.cdist, but writing it by hand shows how broadcasting works.

The dangerous case: silent broadcasting

The most famous broadcasting bug happens in loss functions.

pred = torch.rand(8, 1)    # model output, shape (8, 1)
target = torch.rand(8)     # labels, shape (8,)

wrong = (pred - target)
print(wrong.shape)         # torch.Size([8, 8])  <- not what you wanted!

right = (pred.squeeze(1) - target)
print(right.shape)         # torch.Size([8])

Subtracting (8, 1) and (8,) does not fail. It broadcasts to an 8x8 matrix, and a mean over it gives a plausible-looking but wrong loss. PyTorch's MSELoss warns about this with a user warning when shapes differ. Always make prediction and target shapes identical.

In-place broadcasting limits

In-place operations such as a += b require the result to keep the shape of a. If broadcasting would enlarge a, PyTorch raises an error.

Why broadcasting is efficient

Broadcasting does not actually make copies of the smaller tensor. PyTorch reuses the same memory and simply reads the stretched dimension with a stride of zero. That means adding a bias vector to a large batch costs no extra memory for the bias, and the operation runs in a single fast kernel. This is why broadcast-based code is both shorter and quicker than writing a Python loop over rows.

The same mechanism is behind many layers you will use. Adding a bias to the output of a linear layer broadcasts a vector of shape (out_features,) across the batch. Batch normalisation subtracts a per-channel mean that broadcasts across batch, height and width. In attention, a mask of shape (1, 1, seq, seq) broadcasts across every item and head in the batch. When you read model code and see an operation between two tensors of different sizes, ask yourself which dimensions are stretched, and your understanding of the model will improve.

To practise, try these on small tensors and predict each result shape before running: (5, 1, 3) + (4, 1), (2, 3) * (3,), and (2, 3) + (2,). The first gives (5, 4, 3), the second (2, 3), and the third raises an error because the trailing dimensions 3 and 2 do not match. Getting the third one wrong in your head is the best way to remember the right-alignment rule.

Checklist

• Write the shapes of both operands one above the other, right aligned.
• Mentally check each column: equal, or one is 1.
• Predict the output shape before running, then print it to confirm.
• Use keepdim=True or unsqueeze to make shapes line up deliberately.
• Assert shapes in loss code: assert pred.shape == target.shape.

The next post covers the tools for changing a tensor's shape: view, reshape, permute and squeeze.