PyTorch torch.gather: indices, shapes, and gradients

torch.gather(input, dim, index) selects tensor values along one axis using integer positions. Its index tensor can choose different positions for different rows or batch members, and the output follows the index tensor’s shape. It retrieves existing values; it does not sum them or sort them automatically.

This guide starts with a two-row matrix, then selects whole feature vectors from a batch of sequences. It also compares related indexing operations, reproduces common errors, and checks what happens to gradients when a source value is selected more than once.

The six examples run independently with Python 3.12.14 and PyTorch 2.15.0+cpu. We verified them with 47 checks, including Python-list calculations and a finite-difference gradient comparison. The official gather reference describes the API; the results below come from our CPU examples.

Choose the axis whose coordinates will change

For a matrix and dim=1, each index value selects a column within the same row. With dim=0, it selects a row within the same column. Trace one output coordinate before trying to interpret the whole tensor:

import torch

x = torch.tensor([[10, 11, 12], [20, 21, 22]])
columns = torch.tensor([[2, 0], [1, 1]], dtype=torch.long)
rows = torch.tensor([[1, 0, 1], [0, 1, 0]], dtype=torch.long)
print('Select columns:', torch.gather(x, 1, columns).tolist())
print('Select rows:', torch.gather(x, 0, rows).tolist())

Select columns: [[12, 10], [21, 21]]
Select rows: [[20, 11, 22], [10, 21, 12]]

In the column example, the first row asks for columns 2 and 0, returning [12, 10]. The second asks for column 1 twice, returning [21, 21]. The values in the index are positions, not the numbers to find in the input.

For the row example, output position (0, 0) uses index value 1 and reads x[1, 0], which is 20. Position (0, 1) uses index value 0 and reads x[0, 1], which is 11. Only the coordinate along dim is replaced; the other coordinate remains attached to the output position.

This is why changing dim while keeping the same index tensor can change both the meaning and validity of a call. Our checks compare both outputs with independently written Python-list lookups. They also confirm that dim=-1 selects the final axis in this matrix, just as dim=1 does.

Read the output shape from index

A smaller index tensor does not automatically repeat across the input. In this example, the input has two rows but the index has one:

import torch

x = torch.tensor([[10, 11, 12], [20, 21, 22]])
index = torch.tensor([[2, 0, 1, 2]], dtype=torch.long)
selected = torch.gather(x, 1, index)
print('One index row:', selected.tolist())
print('Output shape:', tuple(selected.shape))
shared = index.expand(x.shape[0], -1)
print('Explicitly expanded:', torch.gather(x, 1, shared).tolist())
empty = torch.empty((2, 0), dtype=torch.long)
print('Empty output shape:', tuple(torch.gather(x, 1, empty).shape))

One index row: [[12, 10, 11, 12]]
Output shape: (1, 4)
Explicitly expanded: [[12, 10, 11, 12], [22, 20, 21, 22]]
Empty output shape: (2, 0)

The first result contains only the first input row’s selections. There is no implicit second row. Calling expand explicitly supplies the same four column choices for both rows, yielding the two-row result shown above.

The four requested columns also demonstrate that the output can be longer along the selected axis than the source. There are only three source columns, but valid positions can be repeated. Every individual index must still point to an existing column.

For nonempty indices, input and index need the same number of axes. On axes other than dim, the index size cannot exceed the input size. The first result satisfies those rules: its one-row index fits inside the two-row input. A three-row index would fail.

The empty index returns shape (2, 0). Treat that as a separate edge case when testing a pipeline. A successful empty-input call does not prove that the same construction will satisfy the checks used for nonempty indices.

Select complete vectors from each sequence

Suppose features has shape (batch, time, feature), and each batch member needs its own selection of time positions. The position list has shape (batch, selected_time). To retrieve every feature at each chosen position, add the feature axis and expand the positions across it:

import torch

features = torch.arange(12).reshape(2, 3, 2)  # batch, time, feature
positions = torch.tensor([[2, 0], [1, 1]], dtype=torch.long)
index = positions.unsqueeze(-1).expand(-1, -1, features.shape[-1])
selected = torch.gather(features, dim=1, index=index)
print('Index shape:', tuple(index.shape))
print('Selected vectors:', selected.tolist())

Index shape: (2, 2, 2)
Selected vectors: [[[4, 5], [0, 1]], [[8, 9], [8, 9]]]

The first sequence selects position 2, then position 0, producing [[4, 5], [0, 1]]. The second selects position 1 twice, producing [[8, 9], [8, 9]]. The repeated time index across both feature coordinates keeps each vector together.

The resulting index and output both have shape (2, 2, 2). The operation preserves the order requested by positions; it does not restore chronological order or remove duplicates. Our verification compares these vectors with direct Python-list selection.

In a training pipeline, check that positions belongs to the same batch as features. The tensor shapes alone cannot detect indices accidentally paired with another batch. Our PyTorch DataLoader guide covers the loading and batching stage that supplies those inputs.

Use index_select for shared positions

When every row needs the same columns, a one-dimensional index list with index_select expresses that directly. Gather can produce the same result after the shared list is explicitly expanded:

import torch

x = torch.tensor([[10, 11, 12], [20, 21, 22]])
common = torch.tensor([2, 0], dtype=torch.long)
print('Shared columns:', torch.index_select(x, 1, common).tolist())
expanded = common.unsqueeze(0).expand(x.shape[0], -1)
print('Gather agrees:', torch.equal(
    torch.index_select(x, 1, common), torch.gather(x, 1, expanded)
))
scores = torch.tensor([[0.2, 0.8, 0.1], [0.9, 0.3, 0.4]], dtype=torch.float64)
best = scores.argmax(dim=1, keepdim=True)
print('Best indices:', best.tolist())
print('Best scores:', torch.gather(scores, 1, best).tolist())
print('take_along_dim agrees:', torch.equal(
    torch.gather(scores, 1, best), torch.take_along_dim(scores, best, dim=1)
))

Shared columns: [[12, 10], [22, 20]]
Gather agrees: True
Best indices: [[1], [0]]
Best scores: [[0.8], [0.9]]
take_along_dim agrees: True

Both shared-column calls return columns 2 and 0 for every row. The earlier gather example instead selected [2, 0] for one row and [1, 1] for another. That per-row choice is the reason to use a matrix of indices.

The second part selects each row’s highest score. Keeping the axis in argmax gives best shape (2, 1), so it can be passed to gather without adding an axis afterward. The selected scores are 0.8 and 0.9.

take_along_dim is also designed to work with indices from operations such as argmax and argsort. It agrees with gather for this specific example. Do not assume that every accepted input shape or default argument is interchangeable between the APIs.

Reproduce common indexing errors

Use torch.long indices in these examples. Check both the shape and the actual index values before a gather call. These five cases deliberately violate different requirements:

import torch

x = torch.tensor([[10, 11, 12], [20, 21, 22]])
cases = {
    'Floating-point indices': torch.tensor([[0., 1.], [1., 2.]]),
    'Missing index axis': torch.tensor([0, 1]),
    'Negative index value': torch.tensor([[0, -1], [1, 2]]),
    'Index beyond final column': torch.tensor([[0, 3], [1, 2]]),
    'Too many index rows': torch.zeros((3, 1), dtype=torch.long),
}
for name, index in cases.items():
    try:
        torch.gather(x, dim=1, index=index)
    except RuntimeError:
        print(name + ': RuntimeError')

Floating-point indices: RuntimeError
Missing index axis: RuntimeError
Negative index value: RuntimeError
Index beyond final column: RuntimeError
Too many index rows: RuntimeError

The floating-point index uses the wrong dtype. The one-dimensional index is missing the row axis. The negative value and value 3 fall outside the valid column positions 0, 1, and 2. The final case has more rows than the input on an axis that is not being selected.

A negative dim and a negative index value mean different things. The dim=-1 check succeeds because it names the final axis. An index value of -1 does not provide Python-style access to the last column in our gather call; it raises an error.

The code prints RuntimeError rather than copying the detailed messages, whose wording can vary between versions. On GPU, also verify the device placement of the source and index tensors. This article tests CPU behavior and does not claim that CUDA reports errors in the same way or at the same point in execution.

Follow repeated selections through backward

Gather can participate in autograd when the source contains floating-point values that require gradients. Repeating a source position repeats its contribution to the loss. The following weighted sum makes those contributions visible:

import torch

x = torch.tensor([[1., 2., 3.], [4., 5., 6.]],
                 dtype=torch.float64, requires_grad=True)
index = torch.tensor([[2, 2, 0], [1, 1, 1]], dtype=torch.long)
weights = torch.tensor([[1., 2., 4.], [1., 3., 5.]], dtype=torch.float64)
selected = torch.gather(x, 1, index)
loss = (selected * weights).sum()
loss.backward()
print('Selected:', selected.tolist())
print('Weighted loss:', loss.item())
print('Input gradient:', x.grad.tolist())
print('Index requires_grad:', index.requires_grad)

Selected: [[3.0, 3.0, 1.0], [5.0, 5.0, 5.0]]
Weighted loss: 58.0
Input gradient: [[4.0, 0.0, 3.0], [0.0, 9.0, 0.0]]
Index requires_grad: False

In the first row, source column 2 is selected twice with weights 1 and 2, so its gradient is 3. Source column 0 is selected once with weight 4. Column 1 is unused and receives zero. In the second row, column 1 receives all three weights: 1 + 3 + 5 = 9.

The weighted loss is 58. Our tests compare the complete gradient matrix against these independent counts and against central finite differences of a Python scalar loss. This checks accumulation, including zero gradients at unselected positions.

The integer index tensor does not carry gradients. If an index was produced by a discrete choice such as argmax, this example does not establish a differentiable path through that choice. It checks gradients with respect to the selected source values.

Gradient recording is a separate decision from selecting values. Our guide to eval, no_grad, and inference_mode explains the controls used for training and inference. These checks use dense CPU gradients; sparse_grad, CUDA performance, and mixed precision need their own validation.

Check coordinates before scaling up

Write down the source axes, the selected axis, and the intended output shape. Then trace one output value back to its source coordinate. For batched sequence selection, verify that each requested position expands across the feature axis you intend to keep. A small exact-value test can reveal an indexing mistake that a shape-only assertion will miss.