Как отсортировать тензор по первому измерению в pytorch?

Question

У меня есть 2D-тензор, и я хотел бы отсортировать его по первому измерению, как в этом примере:

a = torch.FloatTensor(
   [[5, 5],
    [5, 3],
    [3, 5],
    [6, 4],
    [3, 7]])

И я ожидал такого результата после сортировки:

a = torch.FloatTensor(
   [[3, 5],
    [3, 7],
    [5, 3],
    [5, 5],
    [6, 4]])

Это можно ли это сделать в пыторч? Я знаю, что это можно сделать в numpy, но я хочу сделать это в GPU с помощью torch.

Answer 1

Сортировать по первому столбцу и использовать индексы для сортировки всего массива:

a[a[:, 0].sort()[1]]

Вывод:

tensor([[3., 5.],
        [3., 7.],
        [5., 5.],
        [5., 3.],
        [6., 4.]])

И, если вам действительно нужно, чередование:

b = a[a[:, 1].sort()[1]]
b[b[:, 0].sort()[1]]

Вывод:

tensor([[3., 5.],
        [3., 7.],
        [5., 3.],
        [5., 5.],
        [6., 4.]])

Answer 2

torch.stack(sorted(a, key=lambda a: a[0]))

Результат будет:

tensor([[3., 5.],
        [3., 7.],
        [5., 5.],
        [5., 3.],
        [6., 4.]])