Skip to content

← Batching and the Data Path step 2 of 5

Medium Primitives

Pad, and then say where you padded

Jump to the task

Padding puts variable-length sequences into a rectangular tensor. Keep their lengths or a mask so later operations can distinguish real entries from padding:

rows:    [1, 2, 3]        padded: [1, 2, 3, 0]
         [4, 5]                  [4, 5, 0, 0]
         [6, 7, 8, 9]            [6, 7, 8, 9]

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

In this exercise, one means a real entry and zero means padding. Build the mask from the sequence lengths, not from whether values are nonzero: zero can be a valid observation.

Divide by the real count

The second sequence has sum nine and length two, so its mean is 4.5. Averaging its padded row divides by four and gives 2.25 instead. The mean must exclude padding from the count:

import torch

x = torch.tensor([[4., 5., 0., 0.], [0., 2., 0., 0.]])
mask = torch.tensor([[1., 1., 0., 0.], [1., 1., 0., 0.]])
means = (x * mask).sum(dim=1) / mask.sum(dim=1)
print(means.tolist())  # [4.5, 1.0]

The real zero in the second row counts as an observation. This multiplication works for the finite values and zero padding used here. Multiplication by zero does not remove NaN or infinity; if excluded positions may contain those, select or replace them before summing.

Decide what an empty row should mean

An empty sequence has no mathematical mean. This exercise defines its reported mean as zero. Since its masked sum is zero, clamping the count to at least one implements that convention without dividing by zero. Other applications may instead return NaN, mark the result missing or reject empty sequences. Make that choice explicit.

A mask must be used by the next operation

Creating a mask alone does not change a loss or an attention calculation. Pass it in the form that the receiving API expects. Boolean-mask meanings can differ between APIs.

For attention scores, an additive negative-infinity mask can exclude padded keys before softmax when at least one key remains valid. A fully masked row needs defined handling; a plain softmax over all negative infinities produces NaNs. An embedding's padding_idx controls the padding embedding's behavior, but does not automatically mask later attention or loss operations.

Your task

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

The input is a nonempty list of rows containing finite numbers. Rows may be empty. The positive width is at least the length of every row. Return a dictionary with:

  • batch: each row padded on the right with zeros to width.
  • mask: integer ones at real positions and zeros at padding positions.
  • means: each row's mean over real entries, rounded to four decimal places; use 0.0 for an empty row.

The starter constructs the padded batch and mask, but divides by width instead of each row's real count.