torch.take_along_dim

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

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

Если dim равно None, входной массив рассматривается так, как если бы он был преобразован в одномерный.

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

Примечание

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

Параметры:

• input (Tensor) – входной тензор.
• indices (LongTensor) – индексы в input. Должен иметь тип long.
• dim (int, optional) – измерение, вдоль которого выполняется выбор. По умолчанию: None.

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

out (Tensor, optional) – выходной тензор.

Пример:

>>> 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]])