Numpy.lexsort equivalent in pytorch

username: Deepak_Ghimire
cooked: <p>I am looking for numpy.lexsort equivalent function in pytorch. Thanks.</p>
<p>Example:</p>
<blockquote>
<p>import numpy as  np</p>
</blockquote>
<blockquote>
<p>a = np.array([[18.    25.     3.559  0.   ]<br>
[18.    25.     3.559  0.   ]<br>
[18.    25.     3.559  0.   ]<br>
[18.    25.     3.56   0.09 ]<br>
[15.    25.     3.447  0.2  ]<br>
[14.    25.     3.445  0.58 ]<br>
[15.    25.     3.453  0.   ]<br>
[18.    26.     3.559  0.   ]<br>
[18.    26.     3.558  0.   ]<br>
[18.    26.     3.558  0.   ]])</p>
</blockquote>
<blockquote>
<p>indices = np.lexsort((a[:,2], a[:,1], a[:, 0]))<br>
b=a[indices]</p>
</blockquote>
<blockquote>
<p>print(b)</p>
</blockquote>
<p>Output:</p>
<blockquote>
<p>[[14.    25.     3.445  0.58 ]<br>
[15.    25.     3.447  0.2  ]<br>
[15.    25.     3.453  0.   ]<br>
[18.    25.     3.559  0.   ]<br>
[18.    25.     3.559  0.   ]<br>
[18.    25.     3.559  0.   ]<br>
[18.    25.     3.56   0.09 ]<br>
[18.    26.     3.558  0.   ]<br>
[18.    26.     3.558  0.   ]<br>
[18.    26.     3.559  0.   ]]</p>
</blockquote>

username: L_c_Hoang_Duy
cooked: <p>Yes, the pytorch document not include this</p>

username: L_c_Hoang_Duy
cooked: <p>Did you figure out the solution?</p>

username: wouter
cooked: <p>I happened to run into this, I have implemented a lexsort based on stable sort (undocumented behavior) <a href="https://github.com/wouterkool/attention-learn-to-route/blob/master/utils/lexsort.py" rel="noopener nofollow ugc">here</a>.</p>
<p>Also, I give no guarantees, but I think you can use <code>torch.unique</code>:</p>
<pre><code class="lang-auto">def torch_lexsort(a, dim=-1):
    assert dim == -1  # Transpose if you want differently
    assert a.ndim == 2  # Not sure what is numpy behaviour with &gt; 2 dim
    # To be consistent with numpy, we flip the keys (sort by last row first)
    a_unq, inv = torch.unique(a.flip(0), dim=dim, sorted=True, return_inverse=True)
    return torch.argsort(inv)

# Make random float vector with duplicates to test if it handles floating point well
vals = torch.rand(4)
a = vals[(torch.rand(3, 9) * 4).long()]

print(a)
ind = torch_lexsort(a)
ind_np = torch.from_numpy(np.lexsort(a.numpy()))
print("Torch ind", ind)
print("Numpy ind", ind)

print("Torch result")
print (a[:, ind])
print("Numpy result")
print (a[:, ind_np])
</code></pre>
<p>This gives</p>
<pre><code class="lang-auto">tensor([[0.5500, 0.5346, 0.9881, 0.5500, 0.5500, 0.5346, 0.9881, 0.9881, 0.2490],
        [0.5346, 0.5346, 0.5346, 0.5500, 0.2490, 0.5346, 0.5346, 0.9881, 0.2490],
        [0.9881, 0.5500, 0.5500, 0.9881, 0.2490, 0.5346, 0.2490, 0.5346, 0.2490]])
Torch ind tensor([8, 4, 6, 5, 7, 1, 2, 0, 3])
Numpy ind tensor([8, 4, 6, 5, 7, 1, 2, 0, 3])
Torch result
tensor([[0.2490, 0.5500, 0.9881, 0.5346, 0.9881, 0.5346, 0.9881, 0.5500, 0.5500],
        [0.2490, 0.2490, 0.5346, 0.5346, 0.9881, 0.5346, 0.5346, 0.5346, 0.5500],
        [0.2490, 0.2490, 0.2490, 0.5346, 0.5346, 0.5500, 0.5500, 0.9881, 0.9881]])
Numpy result
tensor([[0.2490, 0.5500, 0.9881, 0.5346, 0.9881, 0.5346, 0.9881, 0.5500, 0.5500],
        [0.2490, 0.2490, 0.5346, 0.5346, 0.9881, 0.5346, 0.5346, 0.5346, 0.5500],
        [0.2490, 0.2490, 0.2490, 0.5346, 0.5346, 0.5500, 0.5500, 0.9881, 0.9881]])```</code></pre>

username: hchau630
cooked: <p>Now that pytorch has stable argsort, I rewrote the linked solution by <a class="mention" href="/u/wouter">@wouter</a> more concisely:</p>
<pre><code class="lang-auto">def lexsort(keys, dim=-1):
    if keys.ndim &lt; 2:
        raise ValueError(f"keys must be at least 2 dimensional, but {keys.ndim=}.")
    if len(keys) == 0:
        raise ValueError(f"Must have at least 1 key, but {len(keys)=}.")
    
    idx = keys[0].argsort(dim=dim, stable=True)
    for k in keys[1:]:
        idx = idx.gather(dim, k.gather(dim, idx).argsort(dim=dim, stable=True))
    
    return idx
</code></pre>
<p>compared to the solution based on torch.unique by <a class="mention" href="/u/wouter">@wouter</a>, this can handle arbitrary dim and also handles tensors with NaNs (there is a torch.unique bug that causes it to fail on inputs with NaNs when dim is not None, see <a href="https://github.com/pytorch/pytorch/issues/76571" class="inline-onebox" rel="noopener nofollow ugc">torch.unique() nondeterministic behavior on nan inputs (on GPU) · Issue #76571 · pytorch/pytorch · GitHub</a> and <a href="https://github.com/pytorch/pytorch/issues/95583" class="inline-onebox" rel="noopener nofollow ugc">torch.unique output incorrect when tensor contains NaNs · Issue #95583 · pytorch/pytorch · GitHub</a>. Until this bug is resolved, I don’t see anyway of getting rid of the for-loop while maintaining correctness on inputs with NaNs)</p>

username: xinming_wang
cooked: <p>Thanks very much. I would like to know whether there is a improved method that we can decrease unnecessary sorting operations.</p>