Pad, and then say where you padded

Sequences in a batch have different lengths and a tensor is rectangular, so the short ones get padded:

[1, 2, 3]        [1, 2, 3, 0]
[4, 5]       ->  [4, 5, 0, 0]
[6, 7, 8, 9]     [6, 7, 8, 9]

That is half the job. The other half is a mask saying which positions are real, because every downstream computation now has to be told to ignore the zeros.

mask  [[1, 1, 1, 0],
       [1, 1, 0, 0],
       [1, 1, 1, 1]]

Why padding without a mask is a silent bug

Take the mean of each row. Without a mask, row 2 divides a sum of 9 by 4 instead of by 2, and the answer is 2.25 rather than 4.5. Nothing raises. The numbers are plausible and every short sequence is biased toward zero, proportionally to how short it is.

The same applies to attention, where padded positions get real attention weight, and to losses, where padding contributes to the average.

The masked mean

(x * mask).sum(dim=1) / mask.sum(dim=1)

Multiply to zero out the padding, sum, and divide by the real count rather than the padded width. The keepdim question from the broadcasting track shows up here too if you want to divide the rows rather than reduce them.

Guard the empty row if one is possible: mask.sum() of zero gives a division by zero, and clamp(min=1) on the denominator is the usual fix since the numerator is zero anyway.

Why zero is a bad padding value and also the usual one

Zero is convenient because multiplying by the mask already produces it, and dangerous because it is a plausible data value. For token ids a dedicated pad_id is standard so an embedding lookup can be told to ignore it (nn.Embedding(..., padding_idx=pad_id)). For attention scores the padding value is -inf, from the indexing track, because the softmax has to send it to zero rather than to exp(0).

Your task

def masked_mean(rows: list, width: int) -> dict

Pad each row with zeros to width, build the mask, and return the padded batch, the mask as integers, and each row’s mean over its real entries.