How can I replace the entries of a tensor with the appropriate ranking of each entry?

Question

suppose that I have the following tensor:

>> i = 3
>> j = 5
>> k = 2
>> sor = torch.randn(i,j,k)

>> sor
Out[20]: 
tensor([[[ 0.5604, -0.9675],
         [-1.0953, -0.5615],
         [ 0.4250, -0.9176],
         [-1.6188, -1.0217],
         [-0.0778,  1.9407]],

        [[-0.1034, -0.7925],
         [-0.2955,  0.8058],
         [-0.5349,  1.1040],
         [ 1.1240,  0.8249],
         [ 0.0827, -1.2471]],

        [[ 0.5924,  0.4777],
         [-2.4640, -1.9527],
         [-0.4519,  0.4788],
         [-0.2308, -0.2368],
         [-1.6786,  0.1360]]])

suppose that for every fixed i and j, I want to compute the numeric rank of elements across k, and replace the elements of the tensor sor with those ranks. For instance, from the example above, I want to change the entry [ 0.5604, -0.9675], which is sor[0,0,:] , into [1, 2], since 0.5604 > -0.9675

Thank you,

Answer 1

I think you are looking for torch.argsort:

torch.argsort(sor, dim=2)

Out[ ]:
tensor([[[1, 0],
         [0, 1],
         [1, 0],
         [0, 1],
         [0, 1]],

        [[1, 0],
         [0, 1],
         [0, 1],
         [1, 0],
         [1, 0]],

        [[1, 0],
         [0, 1],
         [0, 1],
         [1, 0],
         [0, 1]]])