← Architecture & Pretraining · Learning notes

Why is logits[:, targets] the wrong selection?

I wanted one correct-token score per row. Instead, indexing silently gave me a whole matrix.

The intuition

For logits shaped (3, 5) and targets [1, 0, 4], this expression selects all three target columns from every row. The result is (3, 3). The intended coordinates are (0, 1), (1, 0), and (2, 4).

logits[torch.arange(3), targets]  # three scores

The intuition is paired coordinates, not a rectangular slice. Once I drew the three selected cells on paper, advanced indexing stopped looking mysterious.

Related reading · NumPy: Advanced indexing