Broadcasting

1. Line the shapes up from the right; missing axes on the left count as size 1.
2. Each pair of sizes must be equal, or one of them must be 1.
3. Size-1 axes stretch to match. Nothing is copied; the same numbers are reused.

$b$ has shape $(4,)$, padded to $(1, 4)$, then stretched down all 3 rows. Every token gets the same bias. With a batch, $(B, T, d) + (d,)$ works the same way.

A column and a row both stretch: $(3, 1) + (1, 4) \to (3, 4)$. Entry $(i, j)$ is $a_i + b_j$, an “outer sum”.

• $(B, T, d) + (d,) \to (B, T, d)$: a bias or gain.
• $(B, h, T, T) + (T, T) \to (B, h, T, T)$: the causal mask.
• $(B, T, d) \times (B, T, 1) \to (B, T, d)$: one scale per token (norms).
• $(B, T, d) + (T,)$: an error, unless $T = d$ happens to hold.

You want to subtract each row’s maximum. $m$ has shape $(2,)$, which aligns with the last axis: it subtracts $m_j$ from column $j$. With a square matrix the shapes still fit, so nothing complains. Keeping the axis, shape $(2, 1)$, subtracts per row.

Line them up from the right and write each result shape, or “error”.

$(5, 3)$ is padded to $(1, 5, 3)$; the 1s stretch. In the second, the last axes are 3 and 2: neither is 1.

x + bias                                   # (B, T, d) + (d,)
scores + mask                              # (B, h, T, T) + (T, T)
x / x.norm(dim=-1, keepdim=True)           # (B, T, d) / (B, T, 1)
torch.broadcast_shapes((4, 1, 3), (5, 3))  # (4, 5, 3)

• Align shapes from the right; sizes must match or be 1; size 1 stretches.
• Bias $(d,)$, mask $(T, T)$ and per-token scale $(B, T, 1)$ all broadcast.
• $(T,)$ aligns with the last axis. Keep reduced axes (keepdim=True) before broadcasting back.