How to rearrange this tensor?


I am searching around for a while, but still can not solve the problem I have at hand. Could somebody point me to the solution, or give an suggestion to this question?

I have a tensor,

A = [[row1], B = [1, 2, 1]
tensor B indicates which group does each row of A belong to, and showing the group number


How can I rearrange A, such that rows for the same group can be grouped together, like:
A_after_rearrange = [[row1]

If I correctly understood your question, this would work. If you want a more general way, you can easily extend it with a loop. It was not a matter of PyTorch, just Python programming…

import torch

a = torch.tensor([[1,2,3],[4,5,6],[7,8,9]])
b = torch.tensor([1,2,1])

rearranged_a = torch.zeros(a.size())
rearranged_a[:2] = a[b==1]
rearranged_a[-1] = a[b==2]

This gives the correct answer! Thank you so much!

Excuse me, but a following question is:
If there is another quicker way instead of extending using for loop?


a = torch.tensor([[1,2,3],[4,5,6],[7,8,9]])
b = torch.tensor([1,2,1])
_, inds = torch.sort(b)
rearranged_a = a[inds]

Oh, right, I think this may be the best answer.

You misunderstood me: I just meant that you should use a for loop but only if you have to do more than two replacements and just for compactness. I also would say that how vector b is built is rather counterintuitive but I don’t know from where it comes. I can guess some class labels. However, if instead of having b like this you manage to have b = [0,2,1], which gives the indeces where you want to place the rows, there exists a builtin method that does the job. Try, with b = [0,2,1]:

rearranged_a = torch.index_select(a,0,b)

Of course, same solution much more compact. I should have thought!

Thank you for giving the informative suggestions!

sorry, I’m back.

If I am using torch.sort(b) then the internal order of all 1s in b is disturbed. It means the inds=[2, 0, 1], not inds=[0, 2, 1]. Is there a way to get b sorted but not disturb elements internal ordering?

b is not modified by the method sort, inds is a new tensor, which gives the argsort of b. However I don’t have clear your question…

Thank you for your answer.

_, inds = torch.sort(b) will give inds = [2, 0, 1]. If I want inds = [0, 2, 1], which means the ordering between two 1s stay the same, then how should I do?

Look, if I run:

_, inds = torch.sort(b)

I get:

[0, 2, 1]

Are you sure you did not modify b?

What if you sort this tensor:

[1, 2, 3, 4, 4, 1, 2, 3, 2, 4, 3, 1, 3, 1, 2, 4]

What will you get? I am getting:

[11,  0, 13,  5, 14,  6,  1,  8, 12,  7, 10,  2,  9,  4,  3, 15]

Yes you are right, the behavior is not the expected.

Hi, I have manage to get expected sort using numpy.argsort(b). This will give stable sort. Thank you for your reply again.

1 Like