torch.gather

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

Собирает значения вдоль оси, заданной параметром dim.

Для 3-мерного тензора результат задается следующим образом:

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 и index должны иметь одинаковое количество измерений. Также требуется, чтобы index.size(d) <= input.size(d) для всех измерений d != dim. out будет иметь ту же форму, что и index. Обратите внимание, что input и index не выполняют широковещательное преобразование относительно друг друга. Если index пуст, всегда возвращается пустой результат той же формы без дополнительной проверки ошибок.

Параметры:

• input (Tensor) – исходный тензор
• dim (int) – ось, вдоль которой выполняется индексация
• index (LongTensor) – индексы собираемых элементов

Именованные аргументы:

• sparse_grad (bool, optional) – Если True, градиент по отношению к input будет разреженным тензором.
• out (Tensor, optional) – целевой тензор

Пример:

>>> t = torch.tensor([[1, 2], [3, 4]])
>>> torch.gather(t, 1, torch.tensor([[0, 0], [1, 0]]))
tensor([[ 1,  1],
        [ 4,  3]])