Pick one per row

You have a score matrix and, for each row, a set of column indices you want from that row:

src            idx           result
[ 0  1  2  3]  [3 0]         [ 3  0]
[ 4  5  6  7]  [1 2]         [ 5  6]
[ 8  9 10 11]  [0 3]         [ 8 11]

Each row uses different columns, so no slice describes it.

Why the loop is not merely slow

for r in range(rows):
    for c in range(cols):
        out[r, c] = src[r, idx[r, c]]

This dispatches a kernel per element. Not per row: per element, several of them, because src[r, idx[r, c]] is itself an indexing op and the assignment is another. On the small case in the tests that is 49 dispatched operations against gather‘s 1.

Each dispatch is a Python call, a schema lookup, a kernel launch and a synchronisation. The arithmetic is free; the ceremony is everything, and you are paying it once per number.

gather

torch.gather(src, dim, index)

index has the same shape as the output. Along dim it says which element to take; along every other axis it means “stay where you are”. So with dim=1, entry [r, c] of the result is src[r, index[r, c]], which is the loop above with the loop deleted.

torch.take_along_dim(src, index, dim=1) is the same operation with a name that reads better and slightly friendlier broadcasting. Use whichever you find clearer; they dispatch the same work.

Where you will meet it

Every time you have per-row indices: the log-probability of the actual next token, the value of the chosen action, the score at the argmax you computed a line earlier. gather is the answer to all of them.

Your task

def pick(src: torch.Tensor, idx: torch.Tensor) -> torch.Tensor

Return the tensor whose [r, c] entry is src[r, idx[r, c]], in at most two dispatched operations.

The starter is the nested loop. It is correct.