Skip to content

← Broadcasting step 6 of 6

Easy Primitives

One bias, every position

This track has mostly shown right-alignment causing trouble. Here is the case it was designed for, and the reason the rule is what it is.

Activations are (batch, time, channels). A bias is one number per channel, (channels,). Adding them:

(batch, time, channels)
              (channels,)
padded to  (1, 1, channels)
---------------------------
(batch, time, channels)

It just works. No unsqueeze, no keepdim, nothing. The bias stretches across every batch item and every time step, which is precisely what a per-channel bias means.

Why the rule is right-aligned

Because the last axis is conventionally the feature axis, and the leading axes are conventionally batch-like. A parameter is per-feature; the batch axes are the ones you want stretched over. Right-alignment makes the common case free and the uncommon case explicit.

Every nn.Linear, LayerNorm and Conv bias in PyTorch relies on this. It is not a convenience that happens to work; it is the convention the whole library is arranged around.

The corollary

When your per-something vector is not per-last-axis, you must say so. That is the entire content of the earlier problems in this track: a per-row quantity in a (rows, cols) tensor needs (rows, 1), because right alignment would otherwise give it to the columns.

So the rule to carry is not “broadcasting is dangerous”. It is: check which axis your vector is per. If it is the last one, do nothing. If it is not, reshape it so its length sits on the axis it belongs to.

Your task

def add_bias(x: torch.Tensor, bias: torch.Tensor) -> torch.Tensor

x is (batch, time, channels) and bias is (channels,). Add the bias to every position, in at most two dispatched operations.

The starter reshapes the bias to something that does not fit.