Search Tech Journey

Find topics, journeys and posts

back to blog
mlbeginner 120m read

DL S003 · Broadcasting — The Rule That Runs the World

Broadcasting is 30% of why NumPy bugs happen and 30% of why NumPy is fast. In this session we derive the three rules from first principles, work through the classic gotchas, and give you the ability to predict any broadcast shape before you run the code.

🧠SoftwareM01 · Math for DL from scratch· Session 003 of 130 120 min

🎯 Read `x[:, None] - centres[None, :]` and predict — without running the code — that the result has shape `(batch, K, d)`. That's broadcasting fluency.

Series: Deep Learning & LLMs From Scratch — 80 sessions · Session 3 / 80 · Module M01 · ~2 hours

The story

I'm going to make an unpopular claim. Broadcasting is the single most underrated concept in numerical computing. More important than matrix multiplication. More consequential than autograd. Here's why: matrix multiplication has one rule, (m, k) @ (k, n) → (m, n), and if you get it wrong Python screams at you immediately. Broadcasting has three rules, and if you get it wrong Python happily runs your code, silently produces a wrong-shape result, and your model trains on garbage until you notice weeks later.

The scary part isn't the rules — they're short. It's the silence when things go wrong. A (3,) array added to a (3, 3) matrix will work. A (3,) array added to a (3, 1) matrix will work. A (1, 3) added to a (3, 1) will work — and give you back a (3, 3), which is almost never what you wanted. Broadcasting is powerful and dangerous in the same breath.

By the end of this session, you'll be able to look at any NumPy or PyTorch line and know, without running it, what shape falls out. That skill alone will save you hundreds of hours over the next 77 sessions.

You will be able to
  • State the three broadcasting rules from memory and apply them in your head to any pair of shapes.
  • Predict the output shape of a mixed-shape operation without running the code.
  • Use `x[:, None]` and `x[None, :]` idioms deliberately to construct desired broadcasts.
  • Recognise the classic 'invisible bug': a shape-`(N,)` combined with a shape-`(N, 1)` producing an accidental `(N, N)`.
  • Explain in one sentence why broadcasting is 'zero-copy' and therefore fast.

Prerequisites

  • Session 001 — you need .shape, .reshape, [:, None], axis= fluency.
  • Session 002 — matrix multiplication and the shape rule.


1 · What is broadcasting, in one sentence?

Broadcasting is NumPy's rule for making two arrays of different shapes 'compatible' for elementwise operations, by pretending the smaller one is repeated along missing or size-1 axes — without actually copying it in memory.

That's the whole thing. Read it twice.

A cookie stamp on a rolled-out sheet of dough
🌍 Real world
💻 Code world

Let's see it in action:

import numpy as np
X = np.array([[1, 2, 3],
              [4, 5, 6]])                # shape (2, 3)
b = np.array([10, 20, 30])                # shape (3,)
 
X + b
# array([[11, 22, 33],
#        [14, 25, 36]])

Notice: X has shape (2, 3), b has shape (3,). They're not the same shape. But NumPy quietly stretched b to look like (2, 3) — i.e., pretended b was [[10, 20, 30], [10, 20, 30]] — and then added element by element. It didn't actually allocate a new array; it just walked through the addition as if b were duplicated.

Why we care: without broadcasting, you'd have to write

b_repeated = np.tile(b, (2, 1))     # actually allocates a (2, 3) array
X + b_repeated

which wastes memory and is slower. Broadcasting is what makes X @ W1 + b1 legal without any manual reshaping.


2 · The three rules

NumPy's broadcasting algorithm has exactly three rules. Learn them cold.

The three broadcasting rules
  • Rule 1 — Align on the RIGHT. If the two shapes have different numbers of dimensions, pad the SHORTER one with 1s on the LEFT until they match.
  • Rule 2 — Compatible axes. Two shapes are compatible on an axis if the sizes are equal OR one of them is 1.
  • Rule 3 — Stretch size-1 axes. Any axis of size 1 is 'stretched' to match the other array's size on that axis. No memory copy — just conceptual repetition.

If at any axis the sizes are neither equal nor include a 1, you get a shape error. That's it. That's every rule.

2.1 Walking through Rule 1 (right-alignment)

Say we're adding shapes (3,) and (2, 3). Not the same length? Pad the shorter ((3,)) with 1s on the left until it has the same rank:

Original shapes: A: (3,) rank 1 B: (2, 3) rank 2 Padded (Rule 1): A: (1, 3) 1 prepended on the left B: (2, 3)

Now they have the same rank. Move to Rule 2.

2.2 Walking through Rule 2 (compatibility)

Axis-by-axis check:

axis 0: A=1, B=2 compatible (one is 1)axis 1: A=3, B=3 compatible (equal) Compatible

2.3 Walking through Rule 3 (stretching)

Any axis where one side is 1 gets stretched to match the other:

axis 0: A stretches from 1 to 2axis 1: no stretching needed Effective A: (2, 3) conceptually [row, row]Effective B: (2, 3) unchanged Elementwise add. Output shape (2, 3).

That's the mental algorithm. Let's practice.


3 · Worked examples — do these on paper

For each pair of shapes, apply Rules 1–3 and predict the output shape.

3.1 Adding a scalar to a matrix

Shapes: () and (3, 4) Rule 1: pad () to (1, 1) then to (1, 1, ...) actually a 0-d array just becomes size-1 everywhere: (1, 1) (3, 4) Rule 2: 1 & 3 OK, 1 & 4 OKRule 3: stretch to (3, 4) Output shape: (3, 4).

Scalars broadcast against literally anything.

np.zeros((3, 4)) + 5     # shape (3, 4), all fives

3.2 Row vector + matrix

Shapes: (3,) and (2, 3) Rule 1: pad (3,) (1, 3) (1, 3) (2, 3) Rule 2: 1 & 2 OK, 3 & 3 OKRule 3: stretch axis 0 to 2 Output shape: (2, 3). this is X + b from §1

3.3 Column vector + row vector = matrix (the WEIRD one)

Shapes: (3, 1) and (1, 4) Rule 1: same rank, no padding needed (3, 1) (1, 4) Rule 2: 3 & 1 OK (one is 1), 1 & 4 OK (one is 1)Rule 3: axis 0 stretches to 3, axis 1 stretches to 4 Output shape: (3, 4).

This one shocks beginners. A column vector plus a row vector gives you an outer product-like matrix:

a = np.array([[1], [2], [3]])         # shape (3, 1)
b = np.array([[10, 20, 30, 40]])      # shape (1, 4)
a + b
# [[11, 21, 31, 41],
#  [12, 22, 32, 42],
#  [13, 23, 33, 43]]

This is powerful — and dangerous. It's the fingerprint of one of the most common bugs in ML code, which we'll cover in §7.

3.4 The one that FAILS

Shapes: (3, 4) and (2, 4) Rule 1: same rank (3, 4) (2, 4) Rule 2: 3 & 2 NEITHER equal NOR one is 1 ERROR
np.zeros((3, 4)) + np.zeros((2, 4))
# ValueError: operands could not be broadcast together with shapes (3,4) (2,4)

Good. NumPy is protecting you.

Try it

For each pair of shapes, apply the three rules and predict the output shape (or "error"):

  1. (5,) and (3, 5)
  2. (5, 1) and (1, 4)
  3. (2, 3, 4) and (4,)
  4. (2, 3, 4) and (3, 4)
  5. (2, 3, 4) and (2, 1, 4)
  6. (2, 3, 4) and (2, 4) ← trap!
  7. (4, 3) and (2, 3)

Answers:

Reveal
  1. (3, 5)(5,) pads to (1, 5), stretches axis 0 to 3.
  2. (5, 4) — 5-tall column meets 4-wide row → 5×4 grid.
  3. (2, 3, 4)(4,) pads to (1, 1, 4), stretches axes 0 and 1.
  4. (2, 3, 4)(3, 4) pads to (1, 3, 4), stretches axis 0.
  5. (2, 3, 4) — axis 1 of the second array is size 1, stretches to 3.
  6. ERROR. (2, 4) pads to (1, 2, 4). Axis 1 is 3 vs 2 — neither equal nor size 1.
  7. ERROR. Same rank, but axis 0 is 4 vs 2.

If you missed #4 or #5, re-read §2.


4 · Constructing broadcasts on purpose — the [:, None] idiom

You saw this in Session 001 §5.2. Now let's use it deliberately.

4.1 The problem

Say you have a batch of 5 vectors, each of dimension 3, and you want to compute the pairwise distances between every pair of vectors. You want a (5, 5) distance matrix.

Naive Python:

distances = np.zeros((5, 5))
for i in range(5):
    for j in range(5):
        distances[i, j] = np.linalg.norm(X[i] - X[j])

Two nested for-loops. Slow. Ugly. But correct.

4.2 The broadcasting solution

Insight: to compute X[i] - X[j] for every pair, we want an array of shape (5, 5, 3) where entry (i, j) is X[i] - X[j]. That's a subtraction between two arrays of shape (5, 1, 3) and (1, 5, 3) — which broadcast to (5, 5, 3).

X = np.random.randn(5, 3)             # (5, 3)
A = X[:, None, :]                     # (5, 1, 3)   — X reshaped as "5 rows, each a mini-batch of 1"
B = X[None, :, :]                     # (1, 5, 3)   — X reshaped as "1 row of 5 vectors"
diffs = A - B                         # (5, 5, 3) — every pairwise difference
distances = np.linalg.norm(diffs, axis=-1)   # (5, 5) — norms along the last axis

Five lines. No Python loop. On big arrays this is ~100× faster than the loop version.

This exact pattern — "insert a size-1 axis so the arrays broadcast against each other in a specific direction" — is one of the two or three most useful moves in all of NumPy. Once you internalize it, you'll use it every day.

4.3 Even shorter, with np.newaxis

np.newaxis is literally an alias for None:

X[:, np.newaxis, :]       # exactly the same as X[:, None, :]

Some people prefer np.newaxis for readability. I use None because it's shorter. Both are correct.


Let's revisit the little forward pass from Session 002 §6.3:

X  = np.random.randn(32, 784)     # 32 flattened MNIST images
W1 = np.random.randn(784, 128)    # hidden layer weights
b1 = np.random.randn(128)         # hidden layer biases

Compute X @ W1 + b1:

  • X @ W1 has shape (32, 128).
  • b1 has shape (128,).

Apply the broadcasting rules:

Rule 1: pad (128,) (1, 128) (32, 128) (1, 128) Rule 2: 32 & 1 OK, 128 & 128 OKRule 3: stretch axis 0 to 32 Effective b1: (32, 128) the same bias vector added to every rowOutput: (32, 128)

One bias vector, added independently to every example in the batch, with zero manual reshaping. That's what broadcasting is buying us.

Every neural network you'll ever write uses this pattern. Every single one.


6 · Broadcasting is zero-copy (this is why it's fast)

I said earlier that NumPy "pretends" the smaller array is duplicated. It doesn't actually allocate.

Let's prove it:

import sys
b = np.arange(1000, dtype=np.float64)          # shape (1000,), ~8 KB
big = np.zeros((1000, 1000), dtype=np.float64) # shape (1000, 1000), ~8 MB
 
result = big + b                               # broadcast add

If broadcasting were done by making a (1000, 1000) copy of b and then adding, we'd allocate an extra 8 MB temporarily. It doesn't. NumPy walks through the addition with a stride trick: for each (i, j) cell, it looks up big[i, j] and adds b[j], using strides to index into b as if it were 2D. Zero copy.

Implication: you can broadcast a small vector against a giant matrix with no memory penalty. This is why we can add a bias vector to every row of a batch of 10,000 examples without breaking a sweat.


7 · War stories — the bugs broadcasting cost me

War story The accuracy that was always exactly 10%

Session 001 §11 already told this one; let me repeat it here because it's the canonical broadcasting bug.

predictions = logits.argmax(axis=-1)             # shape (100,)
labels      = np.array([[0], [1], [2], ...])     # shape (100, 1)  — someone else's data
accuracy    = (predictions == labels).mean()

Shape-(100,) vs shape-(100, 1):

  • Rule 1: pad (100,) to (1, 100)
  • Now we have (1, 100) and (100, 1)
  • Rule 2: 1 & 100 OK, 100 & 1 OK
  • Rule 3: stretches to (100, 100)

The == produced a (100, 100) boolean matrix — comparing every prediction against every label. About 10% of entries were True by chance on a 10-class problem. The .mean() averaged over 10,000 entries. My accuracy was always plausibly ~0.1.

Fix: labels = labels.squeeze() before the comparison. Always sanity-check shapes when a metric looks stuck.

War story The training loss that went sideways

I was computing per-example loss and averaging:

loss_per_example = (predictions - targets) ** 2      # I wanted shape (batch,)
loss = loss_per_example.mean()

predictions was shape (batch,). targets was shape (batch, 1). Broadcasting produced (batch, batch). My "per-example loss" was a batch × batch matrix. The mean was mathematically valid but was averaging the wrong thing. The loss looked plausible, decreased slowly, and my model refused to converge.

Fix: at the point where I load targets, I now always assert targets.shape == predictions.shape. Cheap. Saves days.

War story The image whose channels multiplied instead of scaled

I was normalizing an image by dividing each channel by its per-channel mean:

img = ...                          # shape (3, 224, 224)  — CHW format
means = img.mean(axis=(1, 2))      # shape (3,)
normalized = img / means           # I wanted (3, 224, 224)

img is (3, 224, 224). means is (3,). Broadcasting rules:

  • Rule 1: pad (3,)(1, 1, 3)
  • Now (3, 224, 224) and (1, 1, 3)
  • Rule 2: 3 vs 1 OK. 224 vs 1 OK. 224 vs 3 → ERROR (neither equal nor size 1).

I got a shape error, thankfully. But if img had been square 3-channel — like (3, 3, 3) — I would have silently divided the wrong dimensions.

Fix: always insert explicit None axes to make your intent unambiguous.

means = img.mean(axis=(1, 2))[:, None, None]   # shape (3, 1, 1) — clearly per-channel
normalized = img / means                        # (3, 224, 224) ✅

Common thread across all three war stories: the bug wasn't Python being wrong, it was Python being too helpful. Broadcasting silently made the "wrong" op work. The fix in every case is the same: be explicit about shapes at boundaries. .squeeze(), .reshape(), [:, None], and defensive asserts.


8 · The one weird rule for matmul

@ (matmul) has its own broadcasting behaviour for high-rank tensors, and it's slightly different from elementwise ops. Worth knowing now so you don't get burned in Session 040.

The rule: for A @ B, the LAST TWO axes are treated as the matrix dimensions; everything before them is broadcast batchwise.

A = np.random.randn(32, 5, 3)     # 32 matrices, each (5, 3)
B = np.random.randn(32, 3, 7)     # 32 matrices, each (3, 7)
C = A @ B                          # shape (32, 5, 7) — matmul per batch

You can also broadcast the batch axes:

A = np.random.randn(32, 5, 3)     # 32 matrices, each (5, 3)
B = np.random.randn(3, 7)         # single (3, 7) matrix
C = A @ B                          # shape (32, 5, 7) — B broadcast across the batch

Every transformer attention op uses this. The Q, K, V tensors have shape (batch, heads, seq, d_head); the Q @ K.T op treats the last two axes as the matrix and broadcasts across batch and heads. That's why you can write attention as a two-liner regardless of batch size or head count.


8b · Broadcasting in 2025 transformer code — five patterns you'll meet weekly

Every time a modern LLM computes attention, applies a bias, or does layer normalization, broadcasting is happening under the hood. Here are the five idioms you'll see over and over in nanoGPT, Llama-style code, and HuggingFace transformers — read them until each one triggers an instant “ah, that's a broadcast”.

8b.1 · Adding positional bias to attention scores

# scores has shape (B, nh, T, T)  — batch, heads, query pos, key pos
# alibi_bias has shape (nh, T, T)  — per-head positional bias
scores = scores + alibi_bias        # broadcasts nh, T, T; batch dim stretches

This is ALiBi (Press, Smith & Lewis, 2022) — a linear bias on attention scores that extrapolates to longer sequences than seen at training time. The broadcast is invisible but critical: adding a shape-(nh, T, T) tensor to a shape-(B, nh, T, T) tensor works because axis 0 pads to 1 and stretches. If you accidentally shape the bias as (T, T) instead of (nh, T, T), it still broadcasts — same result across all heads. That's a subtle model change hidden behind a shape choice.

8b.2 · Causal masking

# mask has shape (T, T), lower-triangular of ones, zeros above
# scores has shape (B, nh, T, T)
scores = scores.masked_fill(mask == 0, float("-inf"))

The (T, T) mask broadcasts against (B, nh, T, T). One 4KB mask, applied to a 4GB score tensor. Zero copy. This is the reason a transformer's causal attention is cheap to enforce.

8b.3 · RMSNorm scale

# x has shape (B, T, d)
# gamma has shape (d,)  — learned per-feature scale
rms = torch.sqrt((x ** 2).mean(-1, keepdim=True) + eps)   # (B, T, 1)
x   = x / rms * gamma                                     # (B, T, d)

Two broadcasts in three lines. x / rms: shape-(B, T, d) divided by (B, T, 1). * gamma: shape-(B, T, d) times (d,) — gamma pads to (1, 1, d), stretches. This is RMSNorm (Zhang & Sennrich, 2019), used in Llama, Gemma, Qwen, DeepSeek — essentially every 2024–2025 LLM. LayerNorm has a nearly identical broadcast shape story with an extra mean subtraction and bias.

8b.4 · RoPE rotation

# q has shape (B, nh, T, hd)
# cos, sin have shape (T, hd)  — precomputed rotation table
cos = cos[None, None, :, :]        # (1, 1, T, hd)
sin = sin[None, None, :, :]        # (1, 1, T, hd)
q_rot = q * cos + rotate_half(q) * sin

The None, None axes are added explicitly so the broadcast is unambiguous. Without them, PyTorch would still broadcast (padding on the left), but writing it explicitly is a discipline that saves you when the tensor gains a fifth axis (e.g., for speculative decoding). This is RoPE (Su et al., 2021), refined by YaRN (2023) and LongRoPE (2024) for 128K–10M context windows.

8b.5 · Grouped-Query Attention repeat_interleave

Llama 3, Gemma 2/3, Qwen 2 all use GQA: fewer KV heads than Q heads. To make the shapes match for the Q @ K.T matmul:

# q has shape (B, nh_q, T, hd)   e.g. nh_q = 32
# k has shape (B, nh_kv, T, hd)  e.g. nh_kv = 8
# We want k to broadcast against q's head dim.
k = k.repeat_interleave(nh_q // nh_kv, dim=1)   # (B, nh_q, T, hd)

This is not a pure broadcast — repeat_interleave really allocates. In performance-critical code (vLLM, TensorRT-LLM), this is replaced by a fused kernel that does the effective broadcast without materializing. Same math, different memory story. GQA (Ainslie et al., 2023) trades a tiny quality loss for a huge KV-cache memory saving — typically 4–8×.

Broadcasting in 2025 — muscle memory

    Further reading:


    9 · The full mental model — a decision tree

    When you see any operation A op B, run this mental checklist:

    1. Is `op` elementwise (+, -, *, /, **, ==, <, >, np.exp, np.log, np.maximum, etc.)? Apply the three broadcasting rules to A.shape and B.shape. Predict output shape. 2. Is `op` matmul (@)? Broadcast the LEADING axes elementwise. Matmul the LAST TWO axes with the (m, k) @ (k, n) (m, n) rule. 3. Is `op` a reduction (.sum, .mean, .max, .argmax, .std, .norm)? Look at the axis= argument. That axis is collapsed. Everything else survives. keepdims=True preserves it as size 1.

    That's the whole game. Master this decision tree and shape errors become 10× easier.


    10 · Diagram — the "stretch" mental model

    Small array (5,) Big array (3, 5) a b c d e 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 Broadcasting mentally repeats the small one: a b c d e stretched (conceptually, not in memory) a b c d e a b c d e Then elementwise add: 1+a 2+b 3+c 4+d 5+e 6+a 7+b 8+c 9+d 10+e 11+a 12+b 13+c 14+d 15+e shape (3, 5)

    Draw this on paper once. It burns in.

    Common misconception
    ✗ What most people think

    "Broadcasting is a convenience — it saves me writing a loop. Under the hood NumPy expands the small array to the big shape and then does the elementwise op."

    ✓ What is actually true

    Broadcasting never materialises the expanded array. It sets the stride of the broadcast axis to 0, so the same memory is re-read for every position along that axis. It is a memory trick, not a loop-writing convenience — which is why it costs no extra memory, and why np.broadcast_to returns a read-only array.

    Why the myth is so sticky

    Because the mental picture of "tile it out, then multiply" produces exactly the right answer every single time, so nothing ever corrects you. The model is numerically perfect and only wrong about cost. It stays harmless right up until the moment you write x[:, None] - c[None, :] on a real batch and the process is OOM-killed — because that line genuinely does materialise, and your model gave you no way to distinguish it from the free case.

    Prove it to yourself

    Watch the stride go to zero, and watch the read-only flag appear:

    import numpy as np
    row = np.arange(4)                      # shape (4,)
    b = np.broadcast_to(row, (1_000_000, 4))
    print(b.shape, b.strides)               # (1000000, 4) (0, 8)  <- stride 0
    print(b.nbytes, row.nbytes)             # 32000000 vs 32 (nbytes LIES; no alloc)
    print(b.base is row)                    # True -> same memory
    b[0, 0] = 1                             # ValueError: read-only
    # and the expensive case, for contrast:
    x = np.zeros((1000, 8)); c = np.zeros((256, 8))
    print((x[:, None, :] - c[None, :, :]).nbytes)   # 16 MB, really allocated
    From first principles
    Start with the question

    Why does broadcasting align shapes from the right, and why is a missing axis treated as size 1 rather than an error? This looks like an arbitrary convention — swap it and everything breaks for a reason.

    1. 1
      In C order, the last axis is the fastest-varying one: it is the axis whose elements sit next to each other in memory.
      forced by · an ndarray is a flat buffer, and some axis has to be the contiguous one
    2. 2
      By universal convention the trailing axes are the "what" — features, channels, embedding dimensions — and the leading axes are the "how many" — batch, time, heads.
      forced by · you want the feature vector of one item contiguous in memory so a single item is one cache-friendly run
    3. 3
      The operation you almost always want is "apply this per-feature thing to every item": one bias vector of length d across a batch of (N, d). That is a rule about the trailing axes, replicated over the leading ones.
      forced by · parameters are per-feature; data is per-item; the whole shape of neural network code is exactly that asymmetry
    4. 4
      So the alignment rule must anchor at the trailing end, and must be able to invent leading axes at will.
      forced by · the semantic content lives on the right, and the replication happens on the left
    5. 5
      Inventing a leading axis of size 1 is the unique safe choice: size 1 is the multiplicative identity for shapes, it stretches to any n at zero stride, and it never silently discards data the way a size-n default would.
      forced by · a size-1 axis is the only one that can be re-read for free without committing to a particular n
    ⇒ Therefore

    Therefore right-alignment plus implicit leading 1s is not a convention — it is the only rule consistent with "features are contiguous and on the right, batch is on the left."

    And note what this predicts: adding a bias should be effortless, while adding a per-sample quantity should require an explicit None. Verify it — np.zeros((5, 3)) + np.ones(3) works instantly, while np.zeros((5, 3)) + np.ones(5) raises, and the fix is np.ones(5)[:, None]. It also predicts why keepdims=True exists at all: X.sum(axis=1) deletes the axis, breaking right-alignment, so the sum can no longer be broadcast back against X. Keeping the axis as size 1 is precisely what preserves the alignment.

    Mental modelRight-align and stretch

    Write the two shapes on paper, right-justified like you are adding numbers in a column. Pad the shorter one with 1s on the left. Now read down each column: if the two entries are equal, fine. If one of them is 1, it stretches. Anything else, and NumPy raises.

    Stretching is free — it is a zero stride, the same bytes read over and over. What is not free is when both operands stretch into a new axis: (N,1,d) against (1,K,d) produces (N,K,d), and that result is genuinely allocated. Free stretch, expensive product.

    • Right-align, left-pad with 1s, then per column: equal ⇒ ok · one is 1 ⇒ stretch · else ⇒ error.
    • A stretched axis has stride 0. No memory is used for it, and it is read-only.
    • The output shape is the elementwise max down the columns. Before running a line, compute it on paper and multiply it out — that number is your allocation.
    • None / np.newaxis inserts a size-1 axis to make columns line up. keepdims=True is the same move for reductions.
    🔔 Fires when you see

    Fire this model the moment you see: operands could not be broadcast together · a [:, None] or [None, :] in someone else's code · keepdims=True · a result that came out (N, N) when you wanted (N,) · an OOM on a line with no obvious allocation · pairwise distances, attention scores, or one-hot construction.

    The tradeoff

    You need all-pairs distances between N points and K centres in d dimensions. Broadcast into an (N, K, d) intermediate, or use the matmul expansion of the squared norm?

    Broadcast: x[:, None, :] - c[None, :, :]
    + you gain reads exactly like the maths, works for any distance (L1, Linf, cosine, anything elementwise), and is impossible to get subtly wrong
    − you pay allocates the full N*K*d intermediate — the memory blows up as the product of all three, and it is memory-bandwidth bound rather than compute bound, so it cannot reach BLAS speed no matter what hardware you run it on
    pick when the intermediate N*K*d comfortably fits in RAM with room to spare, or the metric is not L2 so the algebraic shortcut does not exist
    Algebraic: |x|^2 + |c|^2 - 2 x @ c.T
    + you gain peak intermediate is only N*K, the heavy term is a single GEMM which is the most optimised routine in computing, and it hits many-core and GPU throughput properly
    − you pay L2 only; and it is numerically fragile — catastrophic cancellation when points are close can yield small negative "squared distances", so you must clamp at zero before the square root, and it loses precision badly in fp16
    pick when the metric is Euclidean and N*K*d would not fit, or the operation sits in an inner loop where BLAS throughput is what you need
    Chunk the broadcast
    + you gain keeps the readable form and any metric, while capping peak memory at chunk*K*d; trivially parallel across chunks
    − you pay more code, a chunk size to tune, and per-chunk call overhead that dominates if the chunks are small
    pick when N is the axis that grew (streaming or a large dataset) while K and d stay modest — the common production shape
    What a senior engineer actually does

    Write the broadcast version first, always — it is the one you can verify against the definition. Then, before shipping, do the arithmetic: multiply the output shape out and compare it to your memory budget. That single multiplication is the whole skill, and it is the step people skip.

    When you do switch to the algebraic form, keep the broadcast version as the test oracle on a small input, and clamp with np.maximum(d2, 0) before the square root. The cancellation is not hypothetical — it shows up precisely for the nearest neighbours, which are the answers you actually care about.



    11 · Retention scaffold

    Quick recall · click to reveal
    ★ = stretch question

    One-line summary (write it in your own words): _______________________________

    Spaced review: Redo one worked example from §3 tomorrow. Revisit §8b (the five 2025 patterns) on day 7.

    Next session (S004): Derivatives + chain rule + computation graphs — the mathematical machinery that lets losses flow backwards through the tensors you just learned to shape.

    Sticky note: the three rules — (1) align on the right and left-pad with 1s, (2) axes compatible if equal or one is 1, (3) stretch is conceptual (zero-copy strides).


    Recall — from memory

    1. State the three broadcasting rules.
    1. Align on the right — pad shorter with 1s on the left.
    2. Two axes are compatible if equal, or if one is 1.
    3. Stretch size-1 axes to match. Zero copy — conceptual only.
    2. What is the output shape of adding shape (5, 1) to shape (1, 4)?

    (5, 4). Both axes have a 1, so both stretch.

    3. Why is broadcasting fast (i.e., no slower than working with same-shape arrays)?

    It's zero-copy. NumPy uses stride tricks to walk through the operation as if the smaller array were duplicated, without ever allocating a duplicate.

    4. If predictions is shape (100,) and labels is shape (100, 1), what happens when you compute predictions == labels? What shape do you get?

    The (100,) pads to (1, 100), then broadcasts against (100, 1) to produce a (100, 100) boolean matrix — comparing every prediction against every label. This is the classic "accuracy is always 10%" bug. Fix with .squeeze().

    5. For A @ B where A is (32, 5, 3) and B is (3, 7), what's the output shape? Why does this work?

    (32, 5, 7). Matmul treats the last two axes as the matrix; the leading axes broadcast. B has no batch axis, so it's implicitly (1, 3, 7) which broadcasts to (32, 3, 7), then matmul gives (32, 5, 7).

    Stretch — for one extra hour

    Write a NumPy function pairwise_cosine_similarity(A, B) where A has shape (m, d) and B has shape (n, d), and the output has shape (m, n) giving the cosine similarity between every row of A and every row of B. No Python loops. Use broadcasting + matmul.

    Then benchmark it against a nested-for-loop version on m = n = 1000, d = 128. Report the speedup. (Spoiler: 100–500×.)

    In your own words

    Explain broadcasting in one sentence:


    Spaced-review pointer

    • From S001 — you're using .shape, [:, None], and axis semantics constantly.
    • From S002 — the shape rules for @ combine with today's broadcasting for the full transformer attention op.

    Session 004 will use these skills to visualize computation graphs — where every node is a tensor of a specific shape, every edge is an op, and gradients flow through the graph following, yes, exactly the broadcasting rules in reverse.

    Next-session teaser

    You've mastered making shapes agree. Now we ask the deeper question: how does one number in a big graph of numbers depend on another? Session 004 introduces derivatives, the chain rule, and computation graphs — the mathematical machinery that makes learning possible. If today felt like syntax, tomorrow is where things start to move. We'll compute ∂L/∂w by hand on a 3-node computation graph, then generalize it to any graph, then implement it in ~30 lines of NumPy at the end. That's the seed of every autograd engine ever written, including PyTorch's.

    What to bring back tomorrow — sticky note

    • Diagram: the "stretch a small array up to a big array" picture from §10.
    • Equation (well, rules): the three broadcasting rules.
    • Snippet: A - B where A = X[:, None, :] and B = X[None, :, :] gives you all pairwise differences.

    Previous: ← DL S002 · Vectors, Matrices, and the Geometry of Data · Next: DL S004 · Derivatives, the Chain Rule, and Computation Graphs →