How to do a sort that returns the ranking of the numbers in pytorch

1

Hi, everyone. I am wondering how to do a sort that returns the ranking of the numbers?
For example:
a = torch.tensor([[ 1,2,3],[5,6,1]])
the output of torch.argsort(a,axis=1) is tensor([[ 0,1,2],[2,0,1]]), in which the numbers represent the original positions of those cols. But what I want is tensor([[ 0,1,2],[1,2,0]]), in which the number represent the ranking of the original data in that col . Thanks for your attention and time.

2

Hi Tim!

Think of the indices returned by argsort() as permutations. If I understand
what you are asking, you want the inverses of those permutations. They
can be computed as follows:

>>> import torch
>>> torch.__version__
'2.1.1'
>>> a = torch.tensor ([[ 1,2,3], [5,6,1]])
>>> asrt = torch.argsort (a, dim = 1)
>>> asrt
tensor([[0, 1, 2],
        [2, 0, 1]])
>>> torch.empty(2, 3, dtype = torch.long).scatter_ (1, asrt, torch.arange (3).repeat (2, 1))
tensor([[0, 1, 2],
        [1, 2, 0]])

Best.

K. Frank

2 Likes