Pick columns for each row

Each row of a matrix can request its own list of columns. gather reads those positions without a Python loop over the elements.

import torch

src = torch.tensor([[0.0, 1.0, 2.0, 3.0], [4.0, 5.0, 6.0, 7.0]])
idx = torch.tensor([[3, 0], [1, 2]], dtype=torch.long)
print(torch.gather(src, 1, idx).tolist())  # [[3.0, 0.0], [5.0, 6.0]]

The first row selects columns 3 and 0; the second selects columns 1 and 2. The order and any repeated indices are preserved.

Read the shape rule

For dim=1, the output satisfies:

out[r, c] = src[r, idx[r, c]]

src and idx must have the same rank. In this exercise they have the same number of rows, while the number of requested columns can differ from the source width. The output shape equals idx.shape. Each index must name a valid source column.

A slice selects a shared interval or stride pattern. Gather is useful when rows need different choices, including repeated or reordered columns. For these matching row dimensions, take_along_dim can produce the same values, but its broadcasting and dispatch behavior need not be identical.

Why replace the loop

The starter computes the requested values on the CPU inputs, but repeated indexing and assignment dispatch many tensor operations. Gather expresses the selection as one tensor operation. This reduces Python-level overhead; it does not imply that every dispatch launches a GPU kernel or synchronizes a device. Those costs depend on the backend and operation.

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.