Pytorch argsort ordered, with duplicate elements in the tensor
I have a vector A = [0,1,2,3,0,0,1,1,2,2,3,3]. I need to sort it in an increasing matter such that it is listed in an ordered fashion and from that extract the argsort. To better explain this I need to sort A to such that it returns B = [0,4,5,1,6,7,2,8,9,3,10,11]. However, when I use pyotrch's torch.argsort(A) it returns B = [4,5,0,1,6,7,2,8,9,3,10,11]. 

I'm assuming the algorithm that does so cannot be controlled on my end. Is there anyway to approach this without introducing for loops? Such operation are part of my NN model and will cause performance issues if not done efficiently. Thanks!

Here is a pure PyTorch based solution leveraging broadcasting, torch.unique(), and torch.nonzero(). This would be of great boost particularly for a GPU based implementation/run, which is not possible if we have to switch back to NumPy, argsort it and then transfer back to PyTorch (as suggested in other approaches).

# our input tensor
In [50]: A = torch.tensor([0,1,2,3,0,0,1,1,2,2,3,3])

# construct an intermediate boolean tensor
In [51]: boolean = A[:, None] == torch.unique(A)

In [52]: boolean
Out[52]: 
tensor([[1, 0, 0, 0],
        [0, 1, 0, 0],
        [0, 0, 1, 0],
        [0, 0, 0, 1],
        [1, 0, 0, 0],
        [1, 0, 0, 0],
        [0, 1, 0, 0],
        [0, 1, 0, 0],
        [0, 0, 1, 0],
        [0, 0, 1, 0],
        [0, 0, 0, 1],
        [0, 0, 0, 1]], dtype=torch.uint8)


Once we have this boolean tensor, we can find the desired indices by checking for positions where there is an 1 after transposing the boolean tensor.

That would give us both sorted input and the indices. Since we want only the indices, we can just grab those by indexing for the last column (1 or -1)

In [53]: torch.nonzero(boolean.t())[:, -1]
Out[53]: tensor([ 0,  4,  5,  1,  6,  7,  2,  8,  9,  3, 10, 11])




Here's the result for one more example provided by OP in the comments:

In [55]: A_large = torch.tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8, 9, 9])

In [56]: boolean_large = A_large[:, None] == torch.unique(A_large)

In [57]: torch.nonzero(boolean_large.t())[:, -1]
Out[57]: 
tensor([ 0, 10, 11,  1, 12, 13,  2, 14, 15,  3, 16, 17,  4, 18, 19,  5, 20, 21,
         6, 22, 23,  7, 24, 25,  8, 26, 27,  9, 28, 29])


Note: Unlike with NumPy-based solution proposed in other answers, here we don't have to worry about what kind of sorting algorithm we've to use because we are not using any sorting at all.
