43. Collate Function for Variable-Length Sequence Batching

Problem

Real training batches never arrive as neat, equal-length tensors. A DataLoader pulls a list of individually-tokenized examples of different lengths, and it is the collate_fn's job to turn that list into a single padded batch tensor plus an attention mask, so the model can be told which positions are real tokens and which are padding to ignore in attention and loss.

Given a list of variable-length 1D LongTensors (already-tokenized sequences), pad them on the right to the batch's max length with a given pad_id, and produce a boolean attention mask that is True for real tokens and False for padding.

Source: src/43_collate_variable_length_batch.py

def collate_variable_length(
    sequences: list[torch.Tensor], pad_id: int = 0,
) -> tuple[torch.Tensor, torch.Tensor]: ...

>>> seqs = [torch.tensor([5, 6, 7]), torch.tensor([9, 2])]
>>> padded, mask = collate_variable_length(seqs, pad_id=0)
>>> padded
tensor([[5, 6, 7],
        [9, 2, 0]])
>>> mask
tensor([[ True,  True,  True],
        [ True,  True, False]])

Step-by-Step Approach

1. Find max_len, the length of the longest sequence in the batch -- every row of the output will be padded out to this length.
2. Pre-allocate the full (batch, max_len) output tensor filled entirely with pad_id, using the first sequence's dtype so the result stays int64 rather than silently becoming a float tensor.
3. Pre-allocate a same-shaped boolean mask tensor, initialized to False everywhere (padding is the default; real tokens are the exception you write in).
4. For each sequence, copy it into the corresponding row's first length columns of the padded tensor.
5. Set that same row's first length columns of the mask to True -- everything else in that row stays False from the pre-allocation.
6. Return both tensors together; a model consuming this batch uses mask to zero out attention to padding and to exclude padded positions from the loss.

The key insight is pre-allocating both output tensors up front and writing into fixed-size slices, rather than building each row with per-sequence torch.cat/F.pad calls and stacking them afterward -- it avoids one extra tensor allocation per sequence and keeps the whole function to a single pass over the list.

Reference solution

def collate_variable_length(
    sequences: list[torch.Tensor],
    pad_id: int = 0,
) -> tuple[torch.Tensor, torch.Tensor]:
    batch_size = len(sequences)
    lengths = [seq.shape[0] for seq in sequences]
    max_len = max(lengths)

    # Pre-allocate the full (batch, max_len) tensor filled with pad_id, then
    # copy each sequence into its row -- avoids per-row torch.cat/pad calls
    # and one dtype mismatch surprise (new_full keeps sequences[0]'s dtype).
    padded = sequences[0].new_full((batch_size, max_len), pad_id)
    mask = torch.zeros((batch_size, max_len), dtype=torch.bool)

    for i, seq in enumerate(sequences):
        length = lengths[i]
        padded[i, :length] = seq
        mask[i, :length] = True  # real tokens only; padding stays False

    return padded, mask

Key Functions & Tricks

• tensor.new_full(shape, value) — allocates a new tensor with the same dtype and device as an existing tensor, avoiding an accidental dtype mismatch (e.g. padding an int64 sequence tensor with a Python float).
• torch.zeros(shape, dtype=torch.bool) — a clean, explicit way to build a boolean mask that defaults to False.
• Slice-assignment (padded[i, :length] = seq) — writes into a pre-allocated tensor in place, which is cheaper than building each row separately and concatenating.
• seq.shape[0] — reads a 1D tensor's length without a Python len() call that only works because 1D tensors happen to support it.

How to Recognize This Pattern