PyTorch: embeddings, padding and masked sequence pooling

pythontorchindependent-practice

Saved on this device

Full lesson · 4 learning steps Open any exercise for focused practice

Step 1 of 4 · intro card

PyTorch: embeddings, padding and masked sequence pooling

Concept and contract

Average real tokens without letting padding change the representation.

An embedding maps discrete integer identifiers to learned vectors. Batched sequences often use padding so different lengths share a rectangular tensor.

A padding identifier is structural, not an ordinary observed token. Pooling must therefore count real tokens explicitly. Otherwise the same content can receive a different representation merely because it was padded to a longer batch length.

Mechanism and worked example

The fixture defines a four-row, two-feature embedding with known values and padding id zero.

The fixture defines a four-row, two-feature embedding with known values and padding id zero. It looks up vectors, builds a non-padding mask, sums real-token vectors and divides by the real-token count.

clamp_min one gives an all-padding sequence a zero vector without division by zero. The weights and token identifiers are owned synthetic data, not a language model or external vocabulary.

Complete worked example

python

from bootstrap_lab import torch_setup
torch = torch_setup()

from torch import nn

def pooled_tokens(token_ids):
    ids = torch.tensor(token_ids, dtype=torch.long)
    if ids.ndim != 2 or ids.shape[1] == 0:
        raise ValueError('padded token matrix required')
    embedding = nn.Embedding(4, 2, padding_idx=0)
    with torch.no_grad():
        embedding.weight.copy_(torch.tensor([[0.,0.],[1.,2.],[3.,4.],[5.,6.]]))
    vectors = embedding(ids)
    mask = (ids != 0).unsqueeze(-1)
    total = (vectors * mask).sum(dim=1)
    count = mask.sum(dim=1).clamp_min(1)
    return (total / count).tolist()

Reasoning: The fixture defines a four-row, two-feature embedding with known values and padding id zero. It looks up vectors, builds a non-padding mask, sums real-token vectors and divides by the real-token count.

Counterexamples and common mistakes

The counterexample divides by the full padded width instead of the number of real tokens.

The counterexample divides by the full padded width instead of the number of real tokens. A one-token sequence padded to width three is then shrunk by a factor of three.

padding_idx affects embedding behavior but does not automatically fix an arbitrary mean operation. Out-of-range token ids also require a deliberate input policy rather than accidental table indexing failures.

Counterexample: a runnable framework mistake

python

from bootstrap_lab import torch_setup
torch = torch_setup()

from torch import nn

def pooled_tokens(token_ids):
    ids = torch.tensor(token_ids, dtype=torch.long)
    if ids.ndim != 2 or ids.shape[1] == 0:
        raise ValueError('padded token matrix required')
    embedding = nn.Embedding(4, 2, padding_idx=0)
    with torch.no_grad():
        embedding.weight.copy_(torch.tensor([[0.,0.],[1.,2.],[3.,4.],[5.,6.]]))
    vectors = embedding(ids)
    mask = (ids != 0).unsqueeze(-1)
    total = (vectors * mask).sum(dim=1)
    count = torch.tensor(ids.shape[1])
    return (total / count).tolist()

Reasoning: The counterexample divides by the full padded width instead of the number of real tokens. A one-token sequence padded to width three is then shrunk by a factor of three.

Independent practice: one bounded framework change

Add an explicit validation step rejecting token ids outside zero through three before embedding lookup.

Add an explicit validation step rejecting token ids outside zero through three before embedding lookup. Preserve the existing masking and all-padding result.

This introduces one clear vocabulary boundary without adding tokenization or text downloads. Keep integer identifiers as a discrete input contract, and explain that a production tokenizer and vocabulary version must stay compatible with the saved embedding table.

Step 1: Reject ids outside the declared embedding vocabulary. Step 2: Check the minimum and maximum before lookup.

Practice starter: solve before reading the solution

python

from bootstrap_lab import torch_setup
torch = torch_setup()

from torch import nn

def pooled_tokens(token_ids):
    # Remove this placeholder, then adapt the retained baseline:
    # Reject ids outside the declared embedding vocabulary.
    raise NotImplementedError
    ids = torch.tensor(token_ids, dtype=torch.long)
    if ids.ndim != 2 or ids.shape[1] == 0:
        raise ValueError('padded token matrix required')
    embedding = nn.Embedding(4, 2, padding_idx=0)
    with torch.no_grad():
        embedding.weight.copy_(torch.tensor([[0.,0.],[1.,2.],[3.,4.],[5.,6.]]))
    vectors = embedding(ids)
    mask = (ids != 0).unsqueeze(-1)
    total = (vectors * mask).sum(dim=1)
    count = mask.sum(dim=1).clamp_min(1)
    return (total / count).tolist()

Reasoning: Reject ids outside the declared embedding vocabulary. Check the minimum and maximum before lookup.

Hint 1: Raise ValueError for an invalid id.
Hint 2: Keep masking and zero-length denominator protection.

Solution, tests and trade-offs

Cases compare two real tokens with padding, one real token with more padding and an all-padding sequence.

Cases compare two real tokens with padding, one real token with more padding and an all-padding sequence. Practice adds an out-of-range rejection.

Exact vector means expose a wrong denominator, while the invalid case distinguishes explicit validation from a low-level lookup exception. The suite does not evaluate semantic language quality or claim that these four arbitrary vectors encode real word meaning.

Practice solution and acceptance checks

python

from bootstrap_lab import torch_setup
torch = torch_setup()

from torch import nn

def pooled_tokens(token_ids):
    ids = torch.tensor(token_ids, dtype=torch.long)
    if ids.ndim != 2 or ids.shape[1] == 0:
        raise ValueError('padded token matrix required')
    if ids.numel() and ((ids < 0).any() or (ids > 3).any()):
        raise ValueError('token outside vocabulary')
    embedding = nn.Embedding(4, 2, padding_idx=0)
    with torch.no_grad():
        embedding.weight.copy_(torch.tensor([[0.,0.],[1.,2.],[3.,4.],[5.,6.]]))
    vectors = embedding(ids)
    mask = (ids != 0).unsqueeze(-1)
    total = (vectors * mask).sum(dim=1)
    count = mask.sum(dim=1).clamp_min(1)
    return (total / count).tolist()

Reasoning: Cases compare two real tokens with padding, one real token with more padding and an all-padding sequence. Practice adds an out-of-range rejection.

Interview reasoning and next step

More advanced sequence models may use attention masks, positional encodings and task-specific pooling, but padding semantics remain important.

More advanced sequence models may use attention masks, positional encodings and task-specific pooling, but padding semantics remain important. Keep masks aligned with batch and sequence axes and verify them on mixed lengths.

Save vocabulary and preprocessing versions with model state. Before adopting a large architecture, establish that a tiny masked baseline preserves the intended content and ignores only structural padding.

Optional stretch — after the core task works: Add a versioned tokenizer contract with owned variable-length examples.

Computer practice: https://applydjinn.com/typorithm/python-frameworks-lab-20261105.zip. Install the pinned dependencies in an isolated environment, run verify.py --track torch, then edit the retained exercise baseline and check it with --practice pt-embeddings. The lab uses owned in-memory databases, local framework clients and tiny CPU tensors. It never reads the working ApplyDjinn database, contacts production feeds or downloads models.

Average real tokens without letting padding change the representation.