Sorting and rearranging multi-dimensional tensors

SeoHyeong: <p>Hi, I have this task in my hands where I have two tensors of the same two-dimensional size, let’s say tensor <code>a</code> and tensor <code>b</code>. I’d like to sort a and b (in the first dimension) so that the largest element of tensor <code>b</code> is multiplied by the largest element of tensor <code>a</code>, and the second largest element of tensor <code>b</code> is multiplied by the second largest element of tensor <code>a</code>, and so on. However, I want to do this by keeping the original ordering of tensor <code>a</code>.</p>
<p>For example,</p>
<pre><code class="lang-auto">a = torch.tensor([[1,0,2], [1,2,0]])
b = torch.tensor([[0.5, 0.8, 0.1], [0.3, 0.2, 0.1]])
</code></pre>
<p>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.</p>
<pre><code class="lang-auto">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
</code></pre>
<p>Is there a way of performing this calculation?</p>

suraj.pt: <p>Hi <a class="mention" href="/u/seohyeong">@SeoHyeong</a>, AFAIK there isn’t a direct function for sorting by second tensor. You can loop through your tensors to achieve this. Assuming tensor <code>a</code> contains the indices:</p>
<pre><code class="lang-python">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)
</code></pre>
<p>If <code>a</code> contains values instead, you can <code>take</code> by the indices of the sorted values:<br>
<code>rearrange = lambda x, idx: x.sort().values.take(idx.argsort())</code></p>
<p>Hope this helps!</p>

KFrank: <p>Hi SeoHyeong!</p>
<aside class="quote no-group" data-username="SeoHyeong" data-post="1" data-topic="148340" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/seohyeong/48/39759_2.png" class="avatar"> SeoHyeong:</div>
<blockquote>
<p>I want to do this by keeping the original ordering of tensor <code>a</code>.</p>
</blockquote>
</aside>
<p>Yes.  You can do this by multiplying the two sorted tensors together<br>
(so that largest multiplies largest, and so on), but remembering the<br>
“sorting indices” (the <em>argsort()</em>) of tensor <code>a</code>, and then using them to<br>
“unsort” the product tensor so as to recover the original order of <code>a</code>:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
1.10.2
&gt;&gt;&gt;
&gt;&gt;&gt; a = torch.tensor ([[1,0,2], [1,2,0]])
&gt;&gt;&gt; b = torch.tensor ([[0.5, 0.8, 0.1], [0.3, 0.2, 0.1]])
&gt;&gt;&gt;
&gt;&gt;&gt; asort, ainds = a.sort (dim = 1)
&gt;&gt;&gt; bsort, _ = b.sort (dim = 1)
&gt;&gt;&gt;
&gt;&gt;&gt; product_tensor = asort * bsort
&gt;&gt;&gt; wanted_tensor = product_tensor.gather (1, ainds.argsort (dim = 1))
&gt;&gt;&gt;
&gt;&gt;&gt; wanted_tensor
tensor([[0.5000, 0.0000, 1.6000],
        [0.2000, 0.6000, 0.0000]])
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

SeoHyeong: <p>Thank you for the suggestion!</p>

SeoHyeong: <p>Frank, thank you. This is exactly what I was looking for!</p>

SeoHyeong: <p>Hello, Frank. Thanks again for your suggestion. I have another question related to the previous question.</p>
<p>Let’s say there is another tensor called <code>c=torch.tensor(rand(2, 3))</code>.</p>
<p>If I’m planning to rearrange the value of <code>b</code> in order of <code>a</code>, meaning the index where <code>a</code> has the greatest value is occupied by the greatest value from <code>b</code>, and the index where <code>a</code> has the second greatest value is occupied by the second greatest value from <code>b</code>, and multiply the resulting tensor to my new tensor <code>c</code>. This would be exactly what <code>rearragned_b</code> tensor is from the previous question.</p>
<blockquote>
<p><code>rearranged_b = torch.tensor([[0.5, 0.1, 0.8], [0.2, 0.3, 0.1]])</code></p>
</blockquote>
<p>and the final tensor I want you be as follows:</p>
<pre><code class="lang-auto">wanted_c = c * rearranged_b
</code></pre>
<p>Applying your suggestion, I was wondering if the following works. I’ve checked it returns the <code>rearranged_b</code> tensor I was looking for but would you mind double-checking?</p>
<pre><code class="lang-auto">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
</code></pre>
<p>Thank you <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>

KFrank: <p>Hi SeoHyeong!</p>
<aside class="quote no-group" data-username="SeoHyeong" data-post="6" data-topic="148340" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/seohyeong/48/39759_2.png" class="avatar"> SeoHyeong:</div>
<blockquote>
<p>Applying your suggestion, I was wondering if the following works. I’ve checked it returns the <code>rearranged_b</code> tensor I was looking for but would you mind double-checking?</p>
<pre><code class="lang-auto">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
</code></pre>
</blockquote>
</aside>
<p>Yes, this should work, if I understand your goal correctly:</p>
<p>You want be <code>rearranged_b</code> to have its largest value line up with <code>a</code>’s, in<br>
<code>a</code>’s <em>original</em> order, the second largest values to line up, and so on.  Also,<br>
you do <em>not</em> want <code>c</code> to be rearranged – that is, you don’t care where <code>c</code>’s<br>
largest value falls (and so on).</p>
<p>So the <em>order</em> of the values of <code>a</code> affect <code>wanted_c</code>, but the actual values of<br>
<code>a</code> (other than their order) do not affect <code>wanted_c</code>.</p>
<p>Best.</p>
<p>K. Frank</p>

SeoHyeong: <p>This is very correct! Thank you so much! <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=12" title=":smiley:" class="emoji" alt=":smiley:" loading="lazy" width="20" height="20"></p>