Broadcasting: how a [4,256] table adds to a whole batch

Where broadcasting sneaks in

In the last post, the positional-embedding step ended with a line that should have looked suspicious: token embeddings of shape [8, 4, 256] added to a position table of shape [4, 256]. Those shapes don’t match. Element-wise addition normally demands identical shapes — yet PyTorch runs it without complaint. The quiet mechanism making that legal is broadcasting, and it is worth understanding on its own, because it turns up everywhere once you start reading real model code.

The running example, as always: “The cat sat on the mat.” Picture a batch of 8 such sentences, each 4 tokens long, each token a 256-number vector. That is the [8, 4, 256]. The position table only ever has 4 rows — one per slot — so it is [4, 256]. One small table, added to the whole batch.

Start with a scalar

You already use broadcasting without noticing. Add a single number to a vector:

import torch

torch.tensor([1, 2, 3]) + 10   # -> [11, 12, 13]

The 10 is stretched to [10, 10, 10] so the two line up. That stretch is the whole idea — everything below is the same trick on bigger shapes.

Then a row across a matrix

Step up one dimension. Add a row vector to a matrix:

M   = torch.tensor([[0, 0, 0],
                    [10, 10, 10]])   # [2, 3]
row = torch.tensor([1, 2, 3])          # [3]

M + row   # -> [[1, 2, 3], [11, 12, 13]]

The single row is added to every line of the matrix — reused, stretched downward to cover both rows. Hold onto that picture: one vector, applied across many rows. It is exactly what a position table does across a batch.

The rule, in one line

Here is the entire rule. Line the two shapes up from the right. For each axis, they are compatible if they are equal, or one of them is 1 (a missing axis counts as 1). The size-1 axis is the one that gets stretched.

Equal, or one is 1. If two axes are both greater than 1 and unequal, there is no broadcast — it errors.

Back to the batch

Now the shape that started all this:

tok = torch.nn.Embedding(50257, 256)
pos = torch.nn.Embedding(4, 256)

token_emb = tok(inputs)            # [8, 4, 256]
pos_emb   = pos(torch.arange(4))   # [4, 256]
input_emb = token_emb + pos_emb    # [8, 4, 256]

Align from the right: [8, 4, 256] against [4, 256], which is read as [1, 4, 256]. Features match (256 = 256), tokens match (4 = 4), and the batch axis is 1 versus 8 — so the position table is stretched to all 8 sentences. Every sentence gets the same 4 slot vectors added. This is also what breaks the order-blind tie from the previous post: the two identical ' cat' vectors now receive pos[0] and pos[1], and stop being identical.

Nothing actually gets copied

“Stretch” is a mental model, not a memory operation. PyTorch does not build the expanded [8, 4, 256] version of the position table. It re-reads the same 4 rows using zero-stride tricks under the hood. So adding a [4, 256] table — or even a single [256] bias — to a large batch costs essentially nothing in memory. That is why storing one small table and letting broadcasting apply it everywhere is the default move, not an optimization you reach for.

Broadcast, or error?

The failure case is the one that actually bites, so it is worth memorizing. If two axes are both greater than 1 and unequal, there is nothing to stretch, and you get a shape-mismatch error.

Read the last one carefully, because it is a real trap: [8, 4, 256] + [8, 256] looks like it should work — both start with 8 — but alignment is from the right, so 256 meets 256, then 4 meets 8, and it fails. If you meant “one vector per sentence,” you have to spell it [8, 1, 256] so the token axis is an explicit 1.

The takeaway

Broadcasting is two rules and a memory trick: align from the right, stretch any axis that is 1, and never copy. Once it clicks, half of the shape juggling in model code stops being mysterious — including the small, quiet line that gives every token its place in the sentence.