torch.argsort

torch.argsort(input, dim=-1, descending=False, *, stable=False) → Tensor

Возвращает индексы, которые сортируют тензор по заданному измерению в порядке возрастания значений.

Это второе значение, возвращаемое torch.sort(). Точное описание семантики этого метода см. в документации к нему.

Если stable равно True, алгоритм сортировки становится стабильным и сохраняет порядок эквивалентных элементов. Если False, относительный порядок значений, которые считаются равными при сравнении, не гарантируется. True работает медленнее.

Параметры:

• input (Tensor) – входной тензор.
• dim (int, optional) – измерение, по которому выполняется сортировка
• descending (bool, optional) – задаёт порядок сортировки (по возрастанию или убыванию)

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

stable (bool, optional) – задаёт относительный порядок эквивалентных элементов

Пример:

>>> a = torch.randn(4, 4)
>>> a
tensor([[ 0.0785,  1.5267, -0.8521,  0.4065],
        [ 0.1598,  0.0788, -0.0745, -1.2700],
        [ 1.2208,  1.0722, -0.7064,  1.2564],
        [ 0.0669, -0.2318, -0.8229, -0.9280]])


>>> torch.argsort(a, dim=1)
tensor([[2, 0, 3, 1],
        [3, 2, 1, 0],
        [2, 1, 0, 3],
        [3, 2, 1, 0]])