Rearranging a 3-D array using indices from sorting?
I have a 3-D array of random numbers of size [channels = 3, height = 10, width = 10].


Then I sorted it using sort command from pytorch along the columns and obtained the indices as well.



The corresponding index is shown below:



Now, I would like to return to the original matrix using these indices. I currently use for loops to do this (without considering the batches). The code is:

import torch
torch.manual_seed(1)
ch = 3
h = 10
w = 10
inp_unf = torch.randn(ch,h,w)
inp_sort, indices  = torch.sort(inp_unf,1)
resort = torch.zeros(inp_sort.shape)

for i in range(ch):
    for j in range(inp_sort.shape[1]):
        for k in range (inp_sort.shape[2]):
            temp = inp_sort[i,j,k]
            resort[i,indices[i,j,k],k] = temp


I would like it to be vectorized considering batches as well i.e.input size is [batch, channel, height, width].

Using Tensor.scatter_()

You can directly scatter the sorted tensor back into its original state using the indices provided by sort():

torch.zeros(ch,h,w).scatter_(dim=1, index=indices, src=inp_sort)


The intuition is based on the previous answer below. As scatter() is basically the reverse of gather(), inp_reunf = inp_sort.gather(dim=1, index=reverse_indices) is the same as inp_reunf.scatter_(dim=1, index=indices, src=inp_sort):



Previous answer

Note: while correct, this is probably less performant, as calling the sort() operation a 2nd time.

You need to obtain the sorting "reverse indices", which can be done by "sorting the indices returned by sort()".

In other words, given x_sort, indices = x.sort(),  you have x[indices] -&gt; x_sort ; while what you want is reverse_indices such that x_sort[reverse_indices] -&gt; x. 

This can be obtained as follows: _, reverse_indices = indices.sort().

import torch
torch.manual_seed(1)
ch, h, w = 3, 10, 10
inp_unf = torch.randn(ch,h,w)

inp_sort, indices = inp_unf.sort(dim=1)
_, reverse_indices = indices.sort(dim=1)
inp_reunf = inp_sort.gather(dim=1, index=reverse_indices)

print(torch.equal(inp_unf, inp_reunf))
# True

