NumPy Broadcasting

TL;DR

Broadcasting is NumPy’s set of rules for combining arrays of different shapes elementwise as if they had the same shape, without copying. The rules:

1. Compare shapes from the right, axis by axis.
2. Two axes are compatible if they are equal, or one of them is 1.
3. If one shape has fewer axes, prepend 1s until they match.
4. Missing or length-1 axes are stretched (without copying) to the other operand’s size for the operation.

Once you know the rules, broadcasting feels obvious. Until then, it feels like sorcery — (N, 1) - (1, M) somehow produces an (N, M) distance matrix in one line, and you can’t tell what NumPy did.

The practical payoff: tensor operations that would be a triple loop in plain Python become a single line of NumPy. Pairwise distances, column-wise normalisation, batched outer products, attention scores — all of it reduces to a few broadcast-aware multiplications and reductions.

The picture in your head

NumPy never actually duplicates data when broadcasting. It pretends the smaller array has been “stretched” along its length-1 axes by re-using its strides: a length-1 axis with stride 0 means “every index along this axis reads from offset 0” — i.e., the same value over and over. The hardware cost of “broadcasting a (1, M) row over (N, M)” is identical to indexing the same row N times — no real expansion.

Visually, two compatible shapes:

    a:   (    8, 1, 6, 1)
    b:   (       7, 1, 5)
result:  (    8, 7, 6, 5)

Right-aligned. Each axis is either equal or 1 in one of the inputs; the result takes the max along each axis.

Two incompatible shapes:

    a:   (3, 4)
    b:   (3, 5)        # mismatch on the last axis: 4 vs 5, neither is 1

NumPy raises ValueError: operands could not be broadcast together.

The four rules, worked

Take a + b for various shapes. Right-align, check axis-by-axis.

Case 1 — same shape (no broadcasting needed)

a: (3, 4)
b: (3, 4)
result: (3, 4)

Case 2 — scalar broadcast

a: (3, 4)
b: ()                    # zero-D scalar
result: (3, 4)            # b's value applied to every element

a = np.array([[1, 2, 3], [4, 5, 6]])
a * 10
# array([[10, 20, 30], [40, 50, 60]])

Case 3 — row broadcast

a: (3, 4)
b: (   4)                # one length-4 vector
result: (3, 4)            # b is reused for each row

a = np.zeros((3, 4))
b = np.array([1, 2, 3, 4])
a + b
# array([[1., 2., 3., 4.],
#        [1., 2., 3., 4.],
#        [1., 2., 3., 4.]])

Case 4 — column broadcast

a: (3, 4)
b: (3, 1)                # column vector
result: (3, 4)

a = np.zeros((3, 4))
b = np.array([[10], [20], [30]])
a + b
# array([[10., 10., 10., 10.],
#        [20., 20., 20., 20.],
#        [30., 30., 30., 30.]])

Case 5 — outer product via broadcast

a: (N, 1)
b: (1, M)
result: (N, M)            # both axes stretched

a = np.array([[1], [2], [3]])    # (3, 1)
b = np.array([[10, 20, 30]])     # (1, 3)
a * b
# array([[10, 20, 30],
#        [20, 40, 60],
#        [30, 60, 90]])

This is the outer product, expressed without np.outer.

Case 6 — batched outer

a: (B, N, 1)
b: (B, 1, M)
result: (B, N, M)         # batched per-example outer product

This is the shape every transformer attention computation lives in.

A worked numerical example — pairwise distances

Compute pairwise Euclidean distances between two sets of points.

import numpy as np

A = np.array([[0, 0],
              [1, 0],
              [0, 1]], dtype=np.float32)         # (3, 2): 3 points in 2D

B = np.array([[1, 1],
              [2, 2]], dtype=np.float32)         # (2, 2): 2 points in 2D

# Want: D[i, j] = ||A[i] - B[j]||
# Reshape so broadcasting builds the (3, 2, 2) cube of pairwise differences.
A_exp = A[:, None, :]                           # (3, 1, 2)
B_exp = B[None, :, :]                           # (1, 2, 2)

diff = A_exp - B_exp                             # (3, 2, 2): pairwise differences
sq_dist = (diff ** 2).sum(axis=-1)               # (3, 2): sum over the coord axis
D = np.sqrt(sq_dist)

Walk through the shapes:

• A is (3, 2). Insert a length-1 axis at position 1: (3, 1, 2).
• B is (2, 2). Insert a length-1 axis at position 0: (1, 2, 2).
• Subtract: shapes align right (3, 1, 2) and (1, 2, 2). The 1s stretch to the other’s size: (3, 2, 2).
• (diff ** 2).sum(axis=-1) reduces the coord axis (length 2 of (x, y)), leaving (3, 2).

Numerical result:

A = [[0,0], [1,0], [0,1]]
B = [[1,1], [2,2]]

diff[0,0] = (0-1, 0-1) = (-1, -1)   sq_sum = 2,  D = √2 ≈ 1.414
diff[0,1] = (0-2, 0-2) = (-2, -2)   sq_sum = 8,  D = √8 ≈ 2.828
diff[1,0] = (1-1, 0-1) = ( 0, -1)   sq_sum = 1,  D = 1
diff[1,1] = (1-2, 0-2) = (-1, -2)   sq_sum = 5,  D = √5 ≈ 2.236
diff[2,0] = (0-1, 1-1) = (-1,  0)   sq_sum = 1,  D = 1
diff[2,1] = (0-2, 1-2) = (-2, -1)   sq_sum = 5,  D = √5 ≈ 2.236

D = [[1.414, 2.828],
     [1.000, 2.236],
     [1.000, 2.236]]

Three lines of NumPy. The triple loop in plain Python would be ~100× slower for any meaningful N and M.

Inserting axes — the None / np.newaxis idiom

None in an indexing slot inserts a length-1 axis. Same as np.newaxis.

v = np.arange(4)             # (4,)
v[:, None].shape             # (4, 1)
v[None, :].shape             # (1, 4)
v[None, :, None].shape       # (1, 4, 1)

# Equivalent — explicit reshape
v.reshape(4, 1)
v.reshape(1, 4)

The “make a row vector vs make a column vector” pattern:

row = v[None, :]      # (1, n) — row
col = v[:, None]      # (n, 1) — column

# Outer product without np.outer:
M = col * row         # (n, 1) * (1, n) -> (n, n)

Broadcasting in reductions and ufuncs

Broadcasting interacts with axis= reductions and the rest of the NumPy zoo.

X = np.random.default_rng(0).standard_normal((100, 50))

# Per-feature mean/std (column-wise)
mu = X.mean(axis=0, keepdims=True)        # (1, 50)
sd = X.std(axis=0, keepdims=True)         # (1, 50)
X_norm = (X - mu) / sd                     # broadcasts: (100, 50) - (1, 50)

keepdims=True keeps the reduced axis as length 1, which makes the broadcast back to the original shape work naturally. Without keepdims, you’d get (50,) and would need to reshape.

The same pattern in PyTorch is identical, with the same broadcasting rules.

Common gotchas

• Implicit broadcast that you didn’t intend. (N, 1) - (N,) is not an error — it broadcasts to (N, N). If you wanted elementwise subtraction, both should be (N,) or both (N, 1).
• Shape mismatch error messages. operands could not be broadcast together with shapes (3, 4) (3, 5) — read the shapes from the right, find the first incompatible axis. NumPy’s error message tells you the shapes; the rules tell you what’s wrong.
• Memory blow-up. A subtle one: broadcasting itself doesn’t copy, but the result is fully materialised. (10000, 1) * (1, 10000) is a 100M-element matrix — 800 MB at float64. Use chunking or np.einsum (which can sometimes avoid the intermediate).
• (M,) @ (M,) is a dot product, not a 1D × 1D outer. For outer, use M[:, None] * M[None, :] or np.outer.
• PyTorch follows the same rules. Broadcasting in torch is identical to NumPy. Anything you learn here transfers directly.

Quick reference — common axis insertions

Where broadcasting shows up in real ML code

• Layer norm / batch norm — subtract mean, divide by std with keepdims=True aligning the reduction.
• Attention scores — Q @ K.T followed by additive masks of shape (1, 1, T, T) broadcasting against (B, H, T, T).
• Per-class logits + bias — logits + b[None, :] adds the same bias to every row.
• Image normalisation — (image - mean[None, None, :]) / std[None, None, :] to subtract per-channel mean from (H, W, C).
• Pairwise similarity — (A[:, None, :] * B[None, :, :]).sum(-1) for cosine similarity, with normalisation steps before.
• Reward shaping in RL — apply per-step rewards to per-trajectory values via broadcast over the time axis.

The defensive habit: when you write a broadcast expression, write the shapes as a comment.

diff = A[:, None] - B[None]      # (N, 1, D) - (1, M, D) -> (N, M, D)

Six months later, that comment is the difference between “obvious” and “why does this work?”

Resources

• NumPy docs — Broadcasting — numpy.org — canonical reference with diagrams.
• Python Data Science Handbook — Broadcasting chapter — jakevdp.github.io — Jake VanderPlas, free, excellent.
• From Python to Numpy — labri.fr — section on vectorization shows broadcast patterns in depth.
• PyTorch broadcasting docs — pytorch.org — same rules, different syntax.