torch.take_along_dim

torch.take_along_dim(input, indices, dim, *, out=None) → Tensor

Выбирает значения из input по одномерным индексам из indices вдоль указанной dim.

Функции, возвращающие индексы вдоль размерности, такие как torch.argmax() и torch.argsort(), предназначены для работы с этой функцией. См. примеры ниже.

Примечание

Эта функция аналогична функции NumPy’s take_along_axis. См. также torch.gather().

Параметры:

• input (Тензор) – входной тензор.
• indices (тензор) – индексы в input. Должен иметь тип long.
• dim (int) – размерность для выбора.

Ключевые аргументы:

out (Тензор, необязательно) – выходной тензор.

Пример:

>>> t = torch.tensor([[10, 30, 20], [60, 40, 50]])
>>> max_idx = torch.argmax(t)
>>> torch.take_along_dim(t, max_idx)
tensor([60])
>>> sorted_idx = torch.argsort(t, dim=1)
>>> torch.take_along_dim(t, sorted_idx, dim=1)
tensor([[10, 20, 30],
        [40, 50, 60]])