Shapes are the type system of ML, and nobody checks them for you

Part 2 ended on a signature. A batch of sequences of embedding rows is a tensor of shape [B, T, d]. That line is the closest thing this stack has to a type. This post reads it the way you’d read a function’s signature, then shows what the runtime does when the signature is wrong: usually nothing.

I read the broadcasting rule as torch writes it. Then I wrote four shape bugs on purpose and watched which ones the runtime caught. It caught one.

A shape is a signature

A tensor is an N-dimensional array plus a description of how it sits in memory. [B, T, d] says: B sequences in the batch, T positions in each, d numbers at each position. Read it the way you’d read f(batch: List[List[Vec]]), with one difference that matters later: the runtime never looks at the names, only at the sizes.

The memory description is the part software engineers already know. A tensor of shape [4, 8, 16] has strides (128, 16, 1): move one step along the first axis and you skip 128 numbers. Transpose the first two axes and the shape becomes [8, 4, 16] with strides (16, 128, 1), the same memory read in a different order; data_ptr doesn’t change. That’s why transpose is free and why view on the result refuses, with a message that says so, while reshape copies when it has to. A tensor is a pointer, a shape, and a stride tuple; most operations you’ll call are arithmetic on the last two.

Every operation has a shape signature of its own: what goes in, what comes out, which axis disappears. mean(dim=1) on [B, T, d] returns [B, d]; the T is gone. A linear layer maps [B, T, d] to [B, T, d_out]; only the last axis changes. Writing the signature next to the call is the discipline this post is about; nothing else will write it.

The rule is five lines

Two tensors of different shapes can still be added, and the rule that decides how is broadcasting. torch keeps a readable copy of it in Python, under torch.broadcast_shapes, in _refs/__init__.py:392-474. Stripped of the symbolic-shape machinery, the rule is this:

common_shape = [1] * reduce(max, (len(shape) for shape in shapes))
for arg_idx, shape in enumerate(shapes):
    for idx in range(-1, -1 - len(shape), -1):          # align from the right
        if common_shape[idx] == 1:
            common_shape[idx] = shape[idx]              # 1 stretches to anything
        if shape[idx] == 1:
            continue                                    # ... in either direction
        torch._check(common_shape[idx] == shape[idx],   # otherwise sizes must match
            lambda: f"Attempting to broadcast a dimension of length {shape[idx]} at {idx}!")
return common_shape

Align the shapes from the right. A size of 1 stretches to match the other. Equal sizes pass. Anything else raises. That’s the whole rule. The eager path runs the same rule in C++ and raises with a shorter message, the one the widget below shows. It exists so that you can add a bias of shape [d] to [B, T, d] without a loop or a copy.

The rule has one property that the rest of this post is about: it decides by size, not by meaning. [8, 16] + [16] works because 16 equals 16. [8, 16] + [8] raises, because 8 is not 16 and 8 is not 1. Whether the 16 on the right is a sequence length or a hidden size or a batch, the rule has no idea, and neither will you when it’s wrong. When two of your axes happen to be the same size, it can’t tell them apart either. The data scientist has leaned on this rule in numpy for years; the software engineer is meeting a coercion rule with no declared types. Both should read it as the second thing, not the first.

The bug that only happens at batch size sixteen

The first bug. Attention scores have shape [B, T, T], one T × T matrix per sequence. A padding mask has shape [B, T]: which positions in each sequence are real. Adding the mask to the scores switches padding off. The correct form is scores + mask[:, None, :]: the None makes the mask [B, 1, T], so it stretches across the query axis and masks the keys.

Drop the None. With B = 8 and T = 16, scores + mask raises: 16 is not 8. With B = 16 and T = 16, it runs. At B = 1 it runs and is correct by accident: 1 is what the None would have inserted.

The rule aligns [16, 16] against [16, 16, 16] from the right and reads the mask as [1, B, T]. The mask meant for sequence b lands on query row b of every sequence. On random lengths I measured a third of the positions masked wrong and a fifth of the attention weight on padding. The softmax after it is finite, the loss is finite, and the run continues. No warning, because there is nothing to warn about: the rule was followed.

Move the sliders until the two sizes meet and watch the verdict flip from an error to a silent result. Then pick [B, T, 1]: that mask runs at every size and is wrong at every size, because it switches off queries instead of keys. It’s the loud version: a padded query’s row is all -inf and its softmax is NaN. The rule is the same in all three cases.

The second bug is the classic one. Predictions of shape [256], targets of shape [256, 1], subtracted for a squared error. The rule stretches both to [256, 256], every prediction against every target. The mean of that is a number. It trains.

I trained a one-feature linear regression through it for two thousand steps. The loss fell and settled at 9.7. The fitted slope was 0.000, the intercept was the mean of the targets, and the R² was 0.0. The model learned to ignore its input. That’s the best it can do when the loss compares it with everyone else’s target. The library loss, F.mse_loss, warns about this mismatch. A hand-written ((pred - target) ** 2).mean() doesn’t, and the hand-written form is the one that ends up in notebooks.

The third bug never raises where it happens. On a tensor of shape [16, 16, 64], mean(dim=0) and mean(dim=1) both return [16, 64]. Different numbers. Same shape. Nothing downstream can tell. On [3, 5, 7] the same two calls return [5, 7] and [3, 7], and the first thing that tries to combine them raises. The bug was there in both cases; only the sizes decided whether it was visible.

The fourth is the control. [3, 5, 7] plus [3, 5] fails because 7 is not 5. It was caught for the one reason the others weren’t: the sizes differ.

Be fair about the rule. It’s doing what it was asked to. The alternative exists: a runtime that carries axis names and refuses to add T to B. torch shipped one as a prototype, named tensors, and removed it in 2.13, because every library in the stack would have had to agree on the names, and sizes are the only thing all of them share. The rule is right. It is not a type check, and nothing in the stack is.

Padding is a shape decision

Sequences don’t come in rectangles. Three sequences of lengths 5, 3, and 1 become a [3, 5] tensor only by padding, and for those lengths 40% of the rectangle is padding. Every operation over T then needs to know which cells are real. That knowledge is a second tensor, the mask, and it has to travel with the first.

Which side you pad on is part of the signature. Pad on the right and the last real token sits at index length - 1: 4, 2, 0 for the three above. Pad on the left and it sits at 4 for all of them. Decoders that generate one token at a time prefer that; the position to read is always the last. Either choice works. Mixing them is the bug. A model trained right-padded and served left-padded reads its first real token where it expects its last.

The mask is not optional even when the model doesn’t attend. Mean-pool that [3, 5] batch without it and the length-1 sequence’s vector averages one real row with four of padding. On arange(1, 31).reshape(3, 5, 2) the unmasked mean of that sequence is [25, 26]; the masked mean is [21, 22], its only real row. Part 2’s pooling code divides by the mask sum for this reason. The division is the signature, written out.

Memory is shape times bytes

A tensor’s cost is arithmetic. Activations of shape [32, 2048, 4096] are 268 million numbers: 1.07 GB in fp32, 537 MB in bf16, 268 MB in int8. That’s the bill, and you need nothing about the model to compute it: the shape and the width of one number. It decides whether a batch fits on a card.

Drag the three sizes and switch the format. At the defaults one block costs 77 GB in fp32 and the two square rows are 89% of it; pull T down to 128 and they fall to a third. Those two rows are the next part’s subject.

What each format gives up is also arithmetic. fp32 keeps 23 bits of mantissa. bf16 keeps 7 and the whole fp32 exponent: it writes π as 3.140625 and rounds 1 + 0.004 to 1.0078125, but its range is fp32’s, to within the last rounding step. fp16 keeps 10 bits and tops out at 65,504; a gradient can exceed that. That’s why training moved to bf16 and serving reaches for int8. Each is a decision about which precision a tensor doesn’t need.

The batch axis is the one you control at serving time, and it moves two numbers in opposite directions. I timed a [B, 4096] by [4096, 4096] matmul on my laptop’s CPU. It costs 447 µs per row at B = 1 and 22 µs per row at B = 512: twenty times the throughput. The call itself went from 0.45 ms to 11.4 ms; each request waited twenty-five times longer, plus whatever it spent waiting for the batch to fill. Five repetitions on one laptop: the ratio is the finding, not the digits. Dynamic batching is the knob between those two numbers, and its setting is a product decision wearing a shape.

What the ML engineer watches

Four habits, and the first two are one habit.

Annotate shapes at function boundaries. jaxtyping with beartype turns the signature into a check that runs on every call:

@jaxtyped(typechecker=beartype)
def pool(x: Float[Tensor, "B T d"], mask: Bool[Tensor, "B T"]) -> Float[Tensor, "B d"]:
    m = mask[:, :, None].float()
    return (x * m).sum(1) / m.sum(1)

Reduce the wrong axis inside and it raises on the return, because [T, d] is not [B, d]. Pass the column-shaped target into the squared error and it raises before the rule can stretch anything. Two of the four bugs above, caught at the boundary, with the shape written where the next reader looks. The other two need the next habit.

Test with sizes that can’t collide. The bug at batch size sixteen is invisible to the annotation. When B and T are both 16, the names bind to the same number and the check passes; so does the wrong-axis reduce, because [T, d] and [B, d] are the same [16, 64]. The fix is in the tests, not the code: use B = 3, T = 5, d = 7. Every transposition, wrong axis, and dropped None then produces a size that matches nothing, and the runtime raises. Round numbers in tests are how the bug at batch size sixteen survives to production. That’s why the first two are one habit: the annotation is only as strict as the sizes it sees.

Mask and padding side are one decision. Record the side next to the tokenizer config. Compute the last-token index from the mask, never from T. Reject a batch whose mask disagrees with its padding. The pooling code in Part 2 already divides by the mask sum; make sure yours does.

Price a tensor before allocating it. Shape times bytes, at the batch you’ll serve, in the format you’ll serve it in. The widget above is that multiplication, and you know its answer before a single weight is loaded.

None of this is modelling work either. A shape is the one type this stack never checks: a device mismatch raises, a wrong axis broadcasts. The check is yours to write.