Sorting 2D tensor by pairs, not columnwise

seva100: <p>Let’s say I have a 2D tensor <code>A</code> of shape <code>(N, 2)</code>, and I would like to sort its rows as pairs, not each column separately.<br>
In other words, I would like to find an expression which finds a permutation of rows in <code>A</code>, such that if <code>i &lt; j</code>, then I would like this to be true after sorting:<br>
<code>(A[i, 0] &lt; A[j, 0]) or ((A[i, 0] == A[j, 0]) and (A[i, 1] &lt;= A[j, 1])</code></p>
<p>For example, let’s suppose I have the following tensor:</p>
<pre><code class="lang-auto">a = torch.FloatTensor(
   [[5, 5],
    [5, 3],
    [3, 5],
    [6, 4],
    [3, 7]])
</code></pre>
<p>and after sorting by pairs I would like to have:</p>
<pre><code class="lang-auto">[[3, 5]
 [3, 7],
 [5, 3],
 [5, 5],
 [6, 4]])
</code></pre>
<p>At the same time, <code>torch.sort</code> sorts each column as a separate 1D tensor:</p>
<pre><code class="lang-auto">In[]: A.sort(axis=0)
Out[]: 
torch.return_types.sort(
values=tensor([[3., 3.],
        [3., 4.],
        [5., 5.],
        [5., 5.],
        [6., 7.]]),
indices=tensor([[2, 1],
        [4, 3],
        [0, 0],
        [1, 2],
        [3, 4]]))
</code></pre>
<p>Is there any way to do this in PyTorch?<br>
Thank you.</p>

albanD: <p>Hi,</p>
<p>Do you know the maximum value that can be contained in the Tensor? Can the values be negative?</p>

seva100: <p>Hi <a class="mention" href="/u/alband">@albanD</a>,<br>
Let’s suppose that I know the maximum value and the values cannot be negative (I can always renormalize values in my tensor this way). The values can be of float type.</p>

albanD: <p>If you know the max value as <code>M</code>, then you can do:</p>
<pre><code class="lang-auto">augmented_a = a.select(1, 0) * M + a.select(1, 1)
ind = augmented_a.sort().indices

res = a.index_select(0, ind)
</code></pre>
<p>(I didn’t test it but that should work <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:">)</p>
<p>If you don’t know the max, you can play similar tricks if you have a max number of digit and converting them to strings <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=9" title=":wink:" class="emoji" alt=":wink:"></p>

seva100: <p>Thank you!<br>
This works indeed, but I think it can result in some precision loss in some cases. For example, in my particular case the first column has integer values (of type <code>long</code>) and the second column has floating-point type values (<code>float32</code>). When I construct <code>augmented_a</code>, I get a floating-point type 1D array, and only integers in [-16777216, 16777216] can be represented in float32 without precision loss (according to <a href="https://en.wikipedia.org/wiki/Single-precision_floating-point_format#Precision_limits_on_integer_values" rel="nofollow noopener">Wikipedia</a>). <code>a.select(1, 0) * M</code> can contain large numbers like these if <code>M</code> is large.<br>
Maybe there is a more suitable way?</p>

albanD: <p>Yes, this is only the simple solution in the case where numbers are small.<br>
It has the advantage to be fast and only using tensor-wide operations.</p>
<p>I’m not sure how to do this efficiently though.<br>
If you find an efficient algorithm that mostly uses matrices (from numpy or matlab) that does this, do link it here, I’m sure we can adapt it to pytorch.</p>

seva100: <p>In NumPy I could do this via sorting by the 2nd column then stable sorting by the 1st column:</p>
<pre><code class="lang-auto">In [54]: a = np.array(
    ...:     [[5, 5],
    ...:      [5, 3],
    ...:      [3, 5],
    ...:      [6, 4],
    ...:      [3, 7]])
    ...:      

In [55]: inner_sorting = np.argsort(a[:, 1])    # here it does not matter whether we do it in stable way or not

In [56]: a_inner_sorted = a[inner_sorting]

In [57]: a_inner_sorted
Out[57]: 
array([[5, 3],
       [6, 4],
       [5, 5],
       [3, 5],
       [3, 7]])

In [58]: outer_sorting = np.argsort(a_inner_sorted[:, 0], kind='stable')

In [59]: a_outer_sorted = a_inner_sorted[outer_sorting]

In [60]: a_outer_sorted
Out[60]: 
array([[3, 5],
       [3, 7],
       [5, 3],
       [5, 5],
       [6, 4]])
</code></pre>
<p>However I have not found any stable sort in PyTorch. Maybe there is a way to perform stable sort in PyTorch? If so, this method would work for columns of any type I think.</p>

albanD: <p>Right, unfortunately we don’t have stable sort implemented.<br>
There are already some discussion on adding stable sorts (for topk for example: <a href="https://github.com/pytorch/pytorch/issues/27542" rel="nofollow noopener">https://github.com/pytorch/pytorch/issues/27542</a>).<br>
If this is really a blocker for you, you can open a new issue to add a stable sort explaining your use case.</p>

seva100: <p>Great, I will post an issue on GitHub. I could get by with the simple solution, but it might lead to precision loss, so it’s not suitable for me in all cases.<br>
Thank you!</p>

Wesley_Neill: <p>In case someone comes looking here, the NumPy solution mentioned by <a class="mention" href="/u/seva100">@seva100</a> now works with PyTorch. Just use the <code>stable=True</code> keyword argument.</p>