The value at the argmax

Use one tensor to choose a column and another to supply the value at that column. Each row makes its own choice.

import torch

scores = torch.tensor([[1.0, 9.0, 3.0], [7.0, 2.0, 5.0]])
other = torch.tensor([[10.0, 20.0, 30.0], [70.0, 80.0, 90.0]])
idx = scores.argmax(dim=1, keepdim=True)
print(idx.tolist())                    # [[1], [0]]
print(other.gather(1, idx).tolist())    # [[20.0], [70.0]]

The result uses columns 1 and 0 from the two rows. Taking the maximum of other would produce [30.0, 90.0], which ignores the choices from scores.

Keep an index dimension for gather

For this two-dimensional input, gather needs a two-dimensional index tensor. An index of shape (rows, 1) requests one column from each row, and the gathered result has the same shape as that index.

Using argmax(dim=1) without keepdim=True gives shape (rows,), which gather rejects because its rank differs from the input's. Preserve that axis for gathering, then remove it from the result:

selected = other.gather(1, idx).squeeze(1)
print(selected.tolist())  # [20.0, 70.0]

Squeeze dimension 1 specifically so a batch with one row still returns a one-dimensional tensor. torch.take_along_dim(other, idx, dim=1) also works for these matching shapes.

This pattern can look up a value for a greedily selected class or action. If an action was sampled or supplied externally, use its actual index instead of recomputing an argmax.

Select values using coordinates

A slice selects a regularly spaced region of an array. A gather selects entries using supplied indices. Check whether those indices select whole rows or columns, or a different entry within each row.

For [[10, 11], [20, 21]], choosing column 1 from the first row and column 0 from the second gives [11, 20]. Choosing columns [1, 0] from every row instead gives [[11, 10], [21, 20]]. These operations use the same index values but have different meanings and output shapes.

Scatter writes values to selected destinations

Scatter takes values and destination indices. The operation may replace existing values or combine them using a reduction such as addition.

Starting with [0, 0, 0], scatter-add of values [2, 3] at indices [1, 1] gives [0, 5, 0]. Both values contribute to the repeated destination. Repeated assignment is not equivalent to this sum; follow the operation's rule for repeated indices.

Selecting with a mask differs from replacing with a mask

For values [10, 20, 30] and mask [True, False, True], selecting masked entries gives [10, 30]. Replacing masked entries with zero gives [0, 20, 0]. Selection can change the shape; replacement preserves it. A selection whose size depends on runtime values needs particular care in compiled code with fixed output shapes.

Check whether an update changes the original array

JAX indexed updates return a new array value. In PyTorch, basic slicing can return a view that shares storage, while advanced indexing returns a copy. Modifying a selected copy does not modify the source. Direct indexed assignment is a separate operation, so distinguish changing a selection from assigning into the original.