promptdojo_

Broadcasting and vectorization — step 1 of 7

Broadcasting: the rule behind x + b

Deep-learning code is full of expressions like x @ w + b where the shapes don't literally match — x @ w is (batch, out) and b is just (out,). It works because of broadcasting: the framework automatically reuses the smaller tensor across the missing dimensions. Run the editor — that's all broadcasting is, written out by hand.

The rule, precisely

Compare shapes from the right. Two dimensions are compatible when they're equal, or one of them is 1 (or missing). Size-1 (and missing) dimensions get stretched:

(3, 2) + (2,)    -> ok: (2,) aligns with the last dim   -> (3, 2)
(3, 2) + (3, 1)  -> ok: the 1 stretches across columns  -> (3, 2)
(3, 2) + (3,)    -> ERROR: 3 vs 2 on the last dim

That third line is the one that bites: the shapes look related — there are 3 rows and a 3-vector! — but broadcasting aligns from the right, so a per-row vector must be shaped (3, 1), not (3,). unsqueeze/reshape exist for exactly this.

The sharpest version of the trap is silent success: (3,) + (3, 1) broadcasts to (3, 3) — no error, just a matrix you never asked for flowing downstream. When a generated loss curve is flat-wrong and nothing crashed, audit the shapes feeding the loss first: a broadcast (batch,) vs (batch, 1) mismatch inside a loss is a classic.

Vectorization, one more time

Chapter 35's rule graduates: on GPUs it's not just speed hygiene, it's the entire point. A Python loop over 32 samples runs 32 kernel launches of nothing; one (32, ...)-batched tensor op runs once, in parallel. That's why everything in torch is batched — the batch dimension isn't a convenience, it's the performance model. If generated training code loops over individual samples calling the model each time, it's leaving the hardware idle.

Where AI specifically gets this wrong

  • (batch,) vs (batch, 1). The silent (n, n) broadcast inside losses and metrics. print(pred.shape, target.shape) before the loss line — two seconds, saves an evening.
  • Right-alignment amnesia. Cursor "fixes" a broadcast error by transposing something until it runs. Whether it runs correctly is a shape-reading question, not a does-it-crash question. One more, subtler than a bug: loop-over-samples training code. It's correct and slow — and a sign the generator pattern-matched a tutorial from the pre-batch era.