torch.gather¶
- torch.gather(input, dim, index, *, sparse_grad=False, out=None) Tensor ¶
Gathers values along an axis specified by dim.
For a 3-D tensor the output is specified by:
out[i][j][k] = input[index[i][j][k]][j][k] # if dim == 0 out[i][j][k] = input[i][index[i][j][k]][k] # if dim == 1 out[i][j][k] = input[i][j][index[i][j][k]] # if dim == 2
input
andindex
must have the same number of dimensions. It is also required thatindex.size(d) <= input.size(d)
for all dimensionsd != dim
.out
will have the same shape asindex
. Note thatinput
andindex
do not broadcast against each other.- Parameters:
- Keyword Arguments:
Example:
>>> t = torch.tensor([[1, 2], [3, 4]]) >>> torch.gather(t, 1, torch.tensor([[0, 0], [1, 0]])) tensor([[ 1, 1], [ 4, 3]])