Learn AI Create the Future — with Pranay
Main site

PyTorch Tensors

Shapes and broadcasting

Broadcasting is the rule that lets a small tensor stretch to match a big one — the source of PyTorch's shortest code and its sneakiest bugs.

Read these first

On this page 5
  1. Why it exists
  2. How it works
  3. A real example you have seen
  4. Remember this
  5. What to learn next

One lesson, three depths. Pick the one that fits you today — you can switch any time.

Beginner — No maths. Plain English.

Broadcasting means a small tensor is automatically stretched to match a bigger one, so you can combine them without writing a loop.

Think of a teacher announcing "everyone gets 5 bonus marks." She does not walk to each student and say it separately. One sentence, spoken once, applies to the whole class. Broadcasting is PyTorch doing exactly that with numbers.

Why it exists

Say you have marks for 3 students in 2 subjects — a grid. You want to add a different bonus per subject.

Without broadcasting, you would loop over students and add the bonus row each time. Loops in Python are slow and wordy. With broadcasting, you write one plus sign, and PyTorch stretches the bonus row across every student for you.

How it works

The shape of a tensor is the list of its sizes, one per direction: a 3-by-2 grid has shape (3, 2).

When two shapes meet, PyTorch compares them from the right end:

marks   (3, 2)        marks   (3, 2)
bonus      (2)   ✓    wrong      (3)   ✗
         ----                   ----
match sizes from the right:
equal, or one of them is 1, or one is missing  →  stretch to fit
anything else                                  →  error

A size of 1 gets stretched. A missing direction counts as 1. Everything else must match exactly.

A real example you have seen

Every photo filter does this. A photo is a grid of pixels, and "make it warmer" is a tiny tint of three numbers — one per colour channel. The app adds the tiny tint to millions of pixels in one operation. That is broadcasting at work.

Remember this

  • Shape is the size list of a tensor, like (3, 2).
  • Broadcasting stretches sizes of 1 (or missing ones) to fit, automatically.
  • Compare shapes from the right — that is the rule people misremember.

What to learn next

Developer — Code and libraries.

Setup

bash
pip install torch

Outputs verified with torch 2.5.1, CPU.

The three broadcasts you will actually use

broadcasting.py
import torch

prices = torch.tensor([100., 200., 300.])   # shape (3,)
discount = torch.tensor(0.9)                # shape () - one number
print(prices * discount)                    # the number is stretched to fit

marks = torch.tensor([[10., 20.],
                      [30., 40.],
                      [50., 60.]])          # (3, 2): 3 students, 2 subjects
bonus = torch.tensor([1., 100.])            # (2,): one bonus per subject
print(marks + bonus)                        # bonus row is copied to every student

col = torch.tensor([[1.], [2.], [3.]])      # (3, 1)
row = torch.tensor([[10., 20., 30.]])       # (1, 3)
both = col + row                            # BOTH get stretched
print(both.shape)
print(both)
Output
tensor([ 90., 180., 270.])
tensor([[ 11., 120.],
        [ 31., 140.],
        [ 51., 160.]])
torch.Size([3, 3])
tensor([[11., 21., 31.],
        [12., 22., 32.],
        [13., 23., 33.]])

The walkthrough

Scalar × vector is the trivial-looking case that teaches the rule: shape () is treated as all-ones and stretches to anything.

marks + bonus lines up (3, 2) against (2,). From the right: 2 matches 2; the 3 has no partner, so the bonus acts as if copied 3 times. No memory is actually copied — the stretch is an accounting trick, which is why it is free.

col + row is the one to stare at. (3, 1) against (1, 3): both have a 1, both get stretched, and out comes (3, 3) — nine numbers from six. Sometimes that grid is exactly what you want, like a distance table. When you wanted plain element-wise addition, it is a disaster that does not error.

When shapes cannot reconcile, the error is at least honest:

python
torch.zeros(3, 4) + torch.zeros(2, 4)
Output
RuntimeError: The size of tensor a (3) must match the size of tensor b (2) at non-singleton dimension 0

Steering the stretch with unsqueeze

unsqueeze(i) inserts a size-1 direction at position i — it tells broadcasting which way to stretch:

python
v = torch.tensor([1., 2., 3.])
print(v.shape, v.unsqueeze(0).shape, v.unsqueeze(1).shape)
Output
torch.Size([3]) torch.Size([1, 3]) torch.Size([3, 1])

A (3,) vector is ambiguous: against a matrix it behaves as a row. unsqueeze(1) turns it into a column on purpose. Writing the intent explicitly is what separates code that works from code that happens to work.

Common mistakes

The silent (3, 1) + (3,) trap. This produces (3, 3), not (3,). It happens constantly with model outputs: predictions come out (batch, 1), targets are (batch,), and the loss quietly averages a full grid of wrong pairs. The loss even goes down. Squeeze or unsqueeze until both sides match exactly before any loss call.

Reading shapes from the left. (5, 3) and (5,) do NOT broadcast — the 5 would need to align with the 3, from the right. People expect "5 matches 5". Compare right to left, always.

Fixing shape errors by adding dimensions until it runs. Every unsqueeze you cannot explain is a bug you have postponed. If you cannot say which real thing each dimension counts — students, subjects, pixels — stop and print .shape until you can.

Broadcasting away your memory. Stretching is free, but the result is materialised: (10000, 1) + (1, 10000) allocates a hundred-million-element tensor.

Try it yourself

Take marks above and give a different bonus to each student instead of each subject, without writing a loop. You will need one unsqueeze. Predict the output before running it.

What to learn next

Researcher — Mathematics and papers.

The formal rule

Two shapes S and T broadcast to shape U when, aligning from the trailing dimension, for every position i: S_i = T_i, or S_i = 1, or T_i = 1 (absent dimensions count as 1). Then U_i = max(S_i, T_i).

Where S_i, T_i are the sizes at aligned position i. The result rank is max(rank(S), rank(T)).

Implementation: stride 0, not copies

A stretched dimension is implemented by setting that dimension's stride to zero — every index along it reads the same memory. You can see this with expand, the explicit form of broadcasting:

python
v = torch.tensor([1., 2., 3.])
e = v.unsqueeze(0).expand(4, 3)
print(e.stride())
Output
(0, 1)

Consequences worth knowing: expanded views cost O(1) memory and time; writing in-place to a stride-0 view is disallowed precisely because one write would alias many logical elements; and gradient flow through a broadcast sums over the stretched dimensions — the adjoint of "copy" is "add". That last fact is why .grad shapes always match the leaf, not the broadcast result, and it is handled by autograd's internal sum_to_size reduction.

repeat, by contrast, materialises real copies — O(n) memory. Reach for it only when a genuine independent copy per row is needed.

Cost model

Broadcast arithmetic is bandwidth-bound like any element-wise op: O(numel(U)) reads/writes with numel(U) = prod(U_i). The danger case from the developer block — outer-product-shaped results — is a complexity change: (n,1) op (1,n) turns O(n) data into O(n²) output.

Lineage

The semantics are inherited verbatim from NumPy (Harris et al., 2020, Array programming with NumPy, Nature 585), which formalised conventions dating to APL and Yorick. PyTorch adopted NumPy semantics in version 0.2 — earlier releases had stricter, partially incompatible rules, which is why very old code sometimes carries defensive expand_as calls.

Named tensors (Chilamkurthy et al., proposal in torch 1.3) attempted to make alignment errors impossible by matching on dimension names rather than positions; the feature remains prototype-status and unused in mainstream code, so positional discipline — comment every tensor's shape — remains the working practice.

What to learn next