Build Aligned Input and Target Batches

A batch groups $B$ examples without joining their token streams. For context length $T=4$, both input IDs and targets have shape $(B,4)$:

With the first two corpus sequences, one batch is

The batch axis indexes independent sequences. Attention may mix positions within a row under its mask; it must never mix row 0 with row 1.

IDs Stay Integer-Valued

Embedding lookup requires integer token IDs. In PyTorch, the batch therefore uses an integer dtype such as torch.long:

The two arrays are shown together because their row and position alignment is the point of the example. Converting IDs to floating-point numbers does not make them model features; it makes them invalid indices for the embedding lookup.

Assert the Batch Contract

Check the cheap invariants before every new data pipeline is trusted:

These assertions cannot prove that the target is shifted correctly. Add a content check while constructing examples:

Flatten Only at the Loss Boundary

The decoder returns logits $(B,T,V)$. PyTorch cross-entropy can consume the same predictions after flattening batch and time:

The reshape does not change which logits row belongs to which target because both tensors use the same row-major order. Transposing one tensor first would silently pair predictions with the wrong targets while retaining valid shapes.

A Valid Shape Is Necessary, Not Sufficient

Before training, decode random batch rows back to tokens and show their targets. Human-readable examples catch off-by-one and boundary mistakes that shape assertions cannot see.