46. Padding Collate Function for Variable-Length Sequences

Problem

Real batches of token IDs (or audio frame features) almost never share one length — a DataLoader hands your collate_fn a list of variable-length 1-D tensors, and you have to turn that into one rectangular batch tensor before it can go through a model. The standard fix is right-padding every sequence up to the batch's longest one with a pad value, plus a same-shape attention mask so the model (or the loss) can tell real tokens from padding and ignore padding in attention/loss computations.

Implement the collate function: pad every sequence in the batch to the max length and build the corresponding mask.

Source: src/46_pad_collate_attention_mask.py

def pad_collate(
    batch: list[torch.Tensor],
    pad_value: int = 0,
) -> tuple[torch.Tensor, torch.Tensor]:
    ...

Examples:
>>> batch = [torch.tensor([1, 2, 3]), torch.tensor([4, 5])]
>>> padded, mask = pad_collate(batch)
>>> padded
tensor([[1, 2, 3],
        [4, 5, 0]])
>>> mask
tensor([[1, 1, 1],
        [1, 1, 0]])

Step-by-Step Approach

1. Record each sequence's length and find max_len across the batch.
2. Preallocate the output tensor with torch.full((batch_size, max_len), pad_value, dtype=...) rather than concatenating incrementally — one allocation instead of many.
3. Preallocate a same-shape mask tensor of zeros.
4. For each sequence, write it into the padded tensor's row starting at column 0 up to its own length, and set the mask's corresponding prefix to 1.
5. Everything past a sequence's real length stays at the initial pad_value / 0, since it was never overwritten.
6. Return the padded tensor and the mask as a pair.

The key insight is that preallocating both output tensors up front and writing each sequence's real prefix in place is both simpler and faster than building the mask separately from scratch (e.g. via torch.arange comparisons) — the mask falls out of exactly the same loop that does the padding.

Reference solution

def pad_collate(batch, pad_value=0):
    lengths = [seq.size(0) for seq in batch]
    max_len = max(lengths)
    batch_size = len(batch)

    # one allocation, O(batch_size * max_len) total writes
    padded = torch.full((batch_size, max_len), pad_value, dtype=batch[0].dtype)
    mask = torch.zeros((batch_size, max_len), dtype=torch.long)

    for i, seq in enumerate(batch):
        length = seq.size(0)
        padded[i, :length] = seq
        mask[i, :length] = 1  # real tokens only; padding stays 0

    return padded, mask

Key Functions & Tricks

• torch.full — preallocates the padded tensor filled with pad_value in one call, avoiding a separate fill step.
• torch.zeros — preallocates the mask; only the "real token" prefix needs to be overwritten to 1.
• Slice assignment (padded[i, :length] = seq) — writes each variable-length sequence into a fixed-width row without a per-element loop.
• batch[0].dtype — preserves the input dtype (e.g. torch.long for token IDs) instead of defaulting to float.

How to Recognize This Pattern