PyTorch · Function

torch.gather

torch.gather selects values from input along a chosen dimension using an index tensor. The result has the same shape as index.

How it works

torch.gather selects elements from input along dim using the positions supplied by index. The result has index's shape, and selection occurs independently at each coordinate outside the chosen dimension.

Syntax

torch.gather(input, dim, index, *, sparse_grad=False, out=None) → Tensor

Parameters

Name	Type	Description
`input`	`Tensor`	Source tensor containing the values to select.
`dim`	`int`	Axis along which selection occurs.
`index`	`LongTensor`	Tensor containing element indices.
`sparse_grad`	`bool, optional`	If true, the gradient with respect to input is sparse.
`out`	`Tensor, optional`	Destination tensor for the result.

Returns

A Tensor containing the gathered values; its shape is the same as index.

Key points

The dim argument chooses the axis along which indexed values are gathered.
The source and index tensors must have equal rank.
Outside the selected axis, every index dimension must not exceed the corresponding input dimension.
The returned tensor follows index's shape.
Use a LongTensor for element positions.

Common mistakes

input and index must already have compatible dimensionality; torch.gather does not broadcast them.
An empty index returns an empty tensor of the same shape without the usual validation.

Example: Gather values along a dimension

import torch

t = torch.tensor([[1, 2], [3, 4]])
index = torch.tensor([[0, 0], [1, 0]])
result = torch.gather(t, 1, index)
print(result)
# tensor([[1, 1],
#         [4, 3]])

The index tensor selects values along dimension 1, producing one output value for every index entry.

When to use it

Per-position tensor lookup

Select per-row or per-batch positions from a tensor when each output position has its own index.

Indexed model outputs

Choose action or class values using indices generated by another operation.

Sparse gradient computation

Request sparse input gradients when dense gradients are undesirable.

Interview questions

QWhat does torch.gather do?

torch.gather takes input, an axis dim, and a LongTensor index. It returns values selected along that axis, with the output shaped like index.

Follow-up: How do input and index interact in shape checking? They must have the same number of dimensions, and index sizes cannot exceed input sizes in dimensions other than dim.

QHow does dim affect element selection?

For a three-dimensional input, dim determines which coordinate is replaced by the corresponding index coordinate. For example, with dim equal to 1, the middle coordinate comes from index.

Follow-up: What shape does the result have? It always has the same shape as index, rather than necessarily preserving input's shape.

QWhat is the purpose of sparse_grad?

sparse_grad=True requests a sparse gradient with respect to input. The option is a keyword argument and defaults to False.

Follow-up: What other optional keyword argument is documented? out specifies a destination tensor for the result.

torch.scatter
torch.index_select

Source: torch.gather — PyTorch 2.13 documentation. Checked 2026-10-16.