torch.gather() usage interpretation

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

Extracts the corresponding value from \(input\) along the axis specified by \(dim\) and the index specified by \(index\).

For a 3D tensor

\(output[i][j][k]=input[index[i][j][k]][j][k] \quad \#\enspace if \enspace dim==0\)

\(output[i][j][k]=input[i][index[i][j][k]][k] \quad \#\enspace if \enspace dim==1\)

\(output[i][j][k]=input[i][j][index[i][j][k]] \quad \#\enspace if \enspace dim==2\)

The number of \(dimensions\) of \(input\) and \(index\) must be the same.The \(shape\) of \(out\) and \(index\) are the same.(Note the difference between \(dimensions\) and \(shape\))

Example

The following two examples are used to explain the specific usage

Example 1

import torchdim = 0_input = torch.tensor([[10, 11, 12],[13, 14, 15],[16, 17, 18]])index = torch.tensor([[0, 1, 2],[1, 2, 0]])output = torch.gather(_input, dim, index)print(output)# tensor([[10, 14, 18],# [13, 17, 12]])

In this example, _input.shape=(3, 3), dimensions=2, where the dimensions of _input and index are the same as 2, and the shapes of output and index are the same as (2, 3).

Because dim=0, the value of each number in index represents the index of dim=0, that is, the "row" dimension, and the index of each number itself specifies the index of other dimensions.For example, [0, 1, 2] in the 0th row of the index represents the 0th, 1st, and 2nd rows, respectively, and the indices of the three numbers themselves in the dim=1 dimension are 0, 1, and 2, that is, the 0th, 1st, and 2nd.List.Therefore, the first number 0 is located in the 0th row in _input, and 0 itself is in the 0th column in the index, so it is located in the 0th column of _input, so the number 10 is found, and 14 is found in the same way.and 18.

The 1st row [1, 2, 0] in the index represents the 1st, 2nd, 0th row and 0th, 1st, and 2nd column in _input respectively, so find the corresponding number in _input [13, 17, 12].

Example 2

import torchdim = 1_input = torch.tensor([[10, 11, 12],[13, 14, 15],[16, 17, 18]])index = torch.tensor([[0, 1],[1, 2],[2, 0]])output = torch.gather(_input, dim, index)print(output)# tensor([[10, 11],# [14, 15],# [18, 16]])

In this example, _input.shape=(3, 3), dimensions=2, where the dimensions of _input and index are the same as 2, and the shapes of output and index are the same as (3, 2).

Because dim=1, the value of each number in index represents the index of dim=1, that is, the "column" dimension, and the index of each number itself specifies the index of other dimensions.For example, [0, 1] in the 0th row of the index represents the 0th and 1st columns, respectively, and the index of these three numbers in the dim=0 dimension is 0, that is, the 0th row.Therefore, the first number 0 is located in the 0th column in _input, and 0 itself is in the 0th row in the index, so it is located in the 0th row of _input, so the number 10 is found, and 11 is found in the same way.

index represents the 1st, 2nd column and 1st row in _input respectively, so find the corresponding number [14, 15] in _input.

index represents the 2nd, 0th column and 2nd row in _input respectively, so find the corresponding number [18, 16] in _input.

Summary

The above example is a two-dimensional case, and the same can be extended to three or more dimensions.To sum up, the value of each number in index itself represents the index of the dimension specified by the parameter dim, and each other dimension is specified by the index of the corresponding dimension of each number in index.

Reference

torch.gather — PyTorch 1.12 documentation

python - What does the gather function do in pytorch in layman terms? - Stack Overflow