Sorting and rearranging multi-dimensional tensors

Hi, I have this task in my hands where I have two tensors of the same two-dimensional size, let’s say tensor a and tensor b. I’d like to sort a and b (in the first dimension) so that the largest element of tensor b is multiplied by the largest element of tensor a, and the second largest element of tensor b is multiplied by the second largest element of tensor a, and so on. However, I want to do this by keeping the original ordering of tensor a.

For example,

a = torch.tensor([[1,0,2], [1,2,0]])
b = torch.tensor([[0.5, 0.8, 0.1], [0.3, 0.2, 0.1]])

Take a look at the first row of tensors a and b, I want 0.8 from tensor b to be mapped with value 2 from tensor a, 0.5 with the value of 1, and 0.1 with the value of 0. Similarly, for the second row, I’d like to map 0.3 to the value of 2, 0.2 to the value of 1, and 0.1 to the value of 0.

a = torch.tensor([[1,0,2], [1,2,0]])
rearranged_b = torch.tensor([[0.5, 0.1, 0.8], [0.2, 0.3, 0.1]])
wanted_tensor = a * b

Is there a way of performing this calculation?

Hi @SeoHyeong, AFAIK there isn’t a direct function for sorting by second tensor. You can loop through your tensors to achieve this. Assuming tensor a contains the indices:

rearrange = lambda x, idx: x.sort().values.take(idx)
l=[]
for i in range(len(b)):
  l.append(b[i], a[i])
rearranged_b = torch.stack(l)

If a contains values instead, you can take by the indices of the sorted values:
rearrange = lambda x, idx: x.sort().values.take(idx.argsort())

Hope this helps!

1 Like

Hi SeoHyeong!

Yes. You can do this by multiplying the two sorted tensors together
(so that largest multiplies largest, and so on), but remembering the
“sorting indices” (the argsort()) of tensor a, and then using them to
“unsort” the product tensor so as to recover the original order of a:

>>> import torch
>>> print (torch.__version__)
1.10.2
>>>
>>> a = torch.tensor ([[1,0,2], [1,2,0]])
>>> b = torch.tensor ([[0.5, 0.8, 0.1], [0.3, 0.2, 0.1]])
>>>
>>> asort, ainds = a.sort (dim = 1)
>>> bsort, _ = b.sort (dim = 1)
>>>
>>> product_tensor = asort * bsort
>>> wanted_tensor = product_tensor.gather (1, ainds.argsort (dim = 1))
>>>
>>> wanted_tensor
tensor([[0.5000, 0.0000, 1.6000],
        [0.2000, 0.6000, 0.0000]])

Best.

K. Frank

2 Likes

Thank you for the suggestion!

Frank, thank you. This is exactly what I was looking for!

Hello, Frank. Thanks again for your suggestion. I have another question related to the previous question.

Let’s say there is another tensor called c=torch.tensor(rand(2, 3)).

If I’m planning to rearrange the value of b in order of a, meaning the index where a has the greatest value is occupied by the greatest value from b, and the index where a has the second greatest value is occupied by the second greatest value from b, and multiply the resulting tensor to my new tensor c. This would be exactly what rearragned_b tensor is from the previous question.

rearranged_b = torch.tensor([[0.5, 0.1, 0.8], [0.2, 0.3, 0.1]])

and the final tensor I want you be as follows:

wanted_c = c * rearranged_b

Applying your suggestion, I was wondering if the following works. I’ve checked it returns the rearranged_b tensor I was looking for but would you mind double-checking?

a = torch.tensor ([[1,0,2], [1,2,0]])
b = torch.tensor ([[0.5, 0.8, 0.1], [0.3, 0.2, 0.1]])
c = torch.rand(2,3)

asort, ainds = a.sort (dim = 1)
bsort, _ = b.sort (dim = 1)

rearranged_b = bsort.gather(1, ainds.argsort(dim=1))
wanted_c = c * rearranged_b

Thank you :slight_smile:

Hi SeoHyeong!

Yes, this should work, if I understand your goal correctly:

You want be rearranged_b to have its largest value line up with a's, in
a's original order, the second largest values to line up, and so on. Also,
you do not want c to be rearranged – that is, you don’t care where c's
largest value falls (and so on).

So the order of the values of a affect wanted_c, but the actual values of
a (other than their order) do not affect wanted_c.

Best.

K. Frank

1 Like

This is very correct! Thank you so much! :smiley: