DEV Community

PyTorch Broadcasting Explained: The 3 Rules (and the Silent Bug That Bites Everyone)

You added two tensors of different shapes and PyTorch didn't complain — it just returned something bigger than you expected. Or you got a RuntimeError about sizes that don't match and no idea which dimension is wrong. Both come down to one feature: broadcasting.

The short answer

PyTorch broadcasting follows exactly three rules: align shapes from the right, each aligned pair of dimensions must be equal or one of them must be 1, and any size-1 dimension is virtually stretched to match the other. If a pair is neither equal nor 1, you get a RuntimeError. The stretch never copies memory — it's implemented with a stride of 0. Learn the three rules once and broadcasting stops surprising you.

The core example

Add a column of three numbers to a row of four:

import torch

a = torch.tensor([[1], [2], [3]])   # shape (3, 1)
b = torch.tensor([[10, 20, 30, 40]]) # shape (1, 4)
c = a + b
print(c.shape)
print(c)
Enter fullscreen mode Exit fullscreen mode

Output:

torch.Size([3, 4])
tensor([[11, 21, 31, 41],
        [12, 22, 32, 42],
        [13, 23, 33, 43]])
Enter fullscreen mode Exit fullscreen mode

No loop, no manual copying. PyTorch stretched the column across four columns and the row down three rows, then added element by element. Critically, the stretch is virtual — no memory is allocated for the repeated values, which is why broadcasting is fast.

The 3 broadcasting rules

  1. Align from the right. Line the shapes up starting at the last dimension. If one tensor has fewer dimensions, pad it with 1s on the left.
  2. Each pair must match. For each aligned pair, the sizes must be equal, or one of them must be 1.
  3. Size-1 dims stretch. Any dimension of size 1 is virtually repeated to match the other tensor along that axis.

Walking the core example through the rules:

        dim0   dim1
a         3      1   ->  stretches to 4
b         1      4   ->  stretches to 3   (b padded to (1, 4))
result    3      4
Enter fullscreen mode Exit fullscreen mode

The output shape takes the max of each aligned pair: (3, 4).

Worked examples

Scalar + matrix. A scalar has shape () and pads to match anything:

m = torch.ones(3, 4)
print((m + 10).shape)   # (3, 4)
Enter fullscreen mode Exit fullscreen mode

Vector + matrix. A (4,) vector aligns with the last dim of a (3, 4) matrix — it's implicitly treated as (1, 4), then stretched down three rows:

v = torch.tensor([1.0, 2.0, 3.0, 4.0])   # (4,)
m = torch.ones(3, 4)                      # (3, 4)
print((m + v).shape)   # (3, 4)
Enter fullscreen mode Exit fullscreen mode

Per-channel image normalization. Subtract a (3, 1, 1) mean from a (3, 224, 224) image — this is standard ImageNet normalization, and the size-1 dims are exactly what broadcasting knows how to stretch:

img = torch.randn(3, 224, 224)
mean = torch.tensor([.485, .456, .406])
mean = mean[:, None, None]   # (3, 1, 1)
out = img - mean             # (3, 224, 224)
print(out.shape)
Enter fullscreen mode Exit fullscreen mode

Why does broadcasting fail with a RuntimeError?

If a pair of aligned dimensions is neither equal nor 1, PyTorch raises a RuntimeError — it's protecting you from a silent mistake:

a = torch.ones(3)
b = torch.ones(5)
a + b   # RuntimeError!
Enter fullscreen mode Exit fullscreen mode
RuntimeError: The size of tensor a (3) must match the size of tensor b (5)
at non-singleton dimension 0
Enter fullscreen mode Exit fullscreen mode

The fix is always the same: reshape one tensor so the mismatched dimension becomes size 1 (with unsqueeze or reshape), then let broadcasting do the rest. If the shapes genuinely shouldn't combine, the error just saved you from a bug.

The silent broadcasting bug (this one has no error)

The real danger isn't the RuntimeError — that message tells you exactly what's wrong. The danger is when broadcasting succeeds silently and hands you a bigger tensor than you meant:

a = torch.randn(3)      # (3,)
b = torch.randn(3, 1)   # (3, 1)
c = a - b
print(c.shape)          # (3, 3)  <-- probably wrong!
Enter fullscreen mode Exit fullscreen mode

You likely wanted an element-wise (3,) result. Instead (3,) was padded to (1, 3), (3, 1) stayed as is, and you got a 3×3 matrix. No error, no warning — just quietly wrong, and it compounds ten lines later. This is the number-one broadcasting bug.

Broadcasting and memory: the stride-0 trick

The stretch is virtual. Internally PyTorch uses a stride of 0 along the stretched dimension, so the same single row or column is reused for every position:

a = torch.ones(1, 4)
e = a.expand(3, 4)     # explicit broadcast
print(e.shape)         # (3, 4)
print(e.stride())      # (0, 1)  <- stride 0!
print(e.numel())       # 12 "virtual" elements
print(a.numel())       # 4 actual elements
Enter fullscreen mode Exit fullscreen mode

The expanded tensor reports 12 elements but only 4 exist in memory. Stride 0 means "don't advance — reuse the same row." That's why broadcasting is fast even when the virtual expanded tensor would be huge.

Common mistakes and gotchas

  • Subtracting (3,) and (3, 1) expecting (3,). You get (3, 3). Make both shapes match exactly for element-wise.
  • Mismatched non-1 dimensions. (3,) + (5,) raises a RuntimeError at the non-singleton dimension.
  • Not knowing which dim stretched. Print .shape before and after every broadcast until the habit is automatic.
  • Forgetting to add size-1 dims. Per-channel ops need an explicit unsqueeze/None to line up the axes.

A three-line habit that catches every broadcasting bug

# 1. Print both shapes
print(a.shape, b.shape)
# 2. Align from the right on paper:  a: 3,1   b: 1,4  -> out 3,4 — is that what you want?
# 3. If not, unsqueeze or reshape to control which dim stretches
Enter fullscreen mode Exit fullscreen mode

Print, align, decide. That habit catches broadcasting bugs before they turn into a wrong result downstream.

Recap

Align from the right, each pair must be equal or one must be 1, and size-1 dims stretch for free. The output takes the max of each pair. A mismatch raises a RuntimeError; the sneakier failure is a silent blow-up that returns the wrong shape — so print your shapes.


This is one idea from my book PyTorch From Ground Up, which builds everything from tensors upward so nothing stays vague. If it helped, you can grab the free PyTorch tensor cheat-sheet here, run every example from the book in the companion notebooks on GitHub, or find the full paperback on Amazon here.

Top comments (0)