Pad Variable-Length Sequences

Put integer sequences into one rectangular array and return a valid-token mask.

Start with the concrete trace below. It shows the state the algorithm must carry as it runs.

Problem trace

Pad Variable-Length Sequences: Allocate one batch-width rectangle filled with pad values, then copy each sequence and mark the identical half-open slice as valid.

Input and goalPut integer sequences into one rectangular array and return a valid-token mask.

Find the rectangular widthThe sequence lengths are 2, 0, and 1, so width=max(2,0,1)=2 and the output shape is (3,2).

row 0: 3fill cursorrow 0: 4row 1: emptyrow 1: emptyrow 2: 9row 2: empty

inputsequences=[[3,4],[],[9]], pad_value=0lengths[2,0,1]arithmeticmax(2,0,1) = 2

Recognize it

Use padding plus a mask when variable-length sequences must enter rectangular tensor operations without treating synthetic pad cells as data.

Keep true

After processing row r, tokens[r,:len(sequence)] equals the original sequence and mask is true on exactly that slice; every remaining cell retains the pad value and false mask.

Reuse it

For ragged-to-dense conversion, derive one shared extent, initialize safe defaults, and update data and validity metadata with the exact same slices.

Pattern: Allocate once, then fill slices.

Simple idea: Use the longest sequence as the width. Fill the token array with the pad value. Copy each sequence into its row and mark the same positions as valid.

from collections.abc import Sequence
import numpy as np

def pad_sequences(
   sequences: Sequence[Sequence[int]], pad_value: int = 0
) -> tuple[np.ndarray, np.ndarray]:
   width = max((len(sequence) for sequence in sequences), default=0)
   tokens = np.full((len(sequences), width), pad_value, dtype=int)
   mask = np.zeros((len(sequences), width), dtype=bool)

   for row, sequence in enumerate(sequences):
      tokens[row, : len(sequence)] = sequence
      mask[row, : len(sequence)] = True
   return tokens, mask

Cost: $O(batch\times width)$ time and space.