The value at the argmax

You have per-row scores over classes, and a second tensor holding something else per class: a confidence, a log-probability, a reward. You want, for each row, the entry of the second tensor at the position where the first is largest.

scores      [[1.0, 9.0, 3.0],      argmax dim=1 -> [1, 0]
             [7.0, 2.0, 5.0]]

other       [[10., 20., 30.],      want -> [20., 70.]
             [70., 80., 90.]]

Note this is not other.max(dim=1). The maximum of other in row 1 is 90. The answer is 70, because the position came from scores.

The shape trap

idx = scores.argmax(dim=1)      # (rows,)
other.gather(1, idx)            # RuntimeError: Index tensor must have
                                # the same number of dimensions as input

gather wants the index to have the output’s shape, and here the output is per row, which means (rows, 1). So:

idx = scores.argmax(dim=1, keepdim=True)     # (rows, 1)
other.gather(1, idx).squeeze(1)              # (rows,)

keepdim=True again, doing the same job it did in the broadcasting problem: keeping a reduced axis around as a length-1 axis so the next operation can line up against it.

torch.take_along_dim(other, idx, dim=1) is the same thing and its name says what is happening.

Where this is the whole computation

The log-probability the model assigned to the token it actually chose. The Q-value of the action taken. The confidence at the predicted class. All of them are “argmax over one tensor, read another at that position”, and all of them are one argmax and one gather.

Your task

def value_at_best(scores: torch.Tensor, other: torch.Tensor) -> torch.Tensor

Return a 1-D tensor of length rows: for each row, the entry of other at the position where scores is largest.