CODEWITHSUNDEEP
pythonsortingpytorchtensortorch
Fernando Noronha
Fernando Noronha

Reputation: 91

How to sort a tensor by first dimension

I have a 2D tensor and I would like to sort by the first dimension like this example:

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

And I expected this result after sorting:

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

Is it possible to do this in pytorch? I know that is possible to do it in numpy, but I want do it in GPU using torch.

Upvotes: 9

Views: 7747

Answers (2)

cosmic_inquiry
cosmic_inquiry

Reputation: 2684

Sort by first column and use the indices to then sort the whole array:

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

Output:

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

And if you really need it interleaved:

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

Output:

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

Upvotes: 15

汉林张
汉林张

Reputation: 19

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

Output will be:

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

Upvotes: 1

Related Questions