How to implement a more general argsort

username: usami
cooked: <p>In <code>torch.sort</code> or <code>torch.argsort</code>, I can specify <code>descending</code> to <code>True</code> or <code>False</code> to get sorted order I want.<br>
How can I specify the sorting by a dynamic comparison order. For example, given following tensor to be sorted:</p>
<pre><code class="lang-auto">a = torch.tensor([[2, 1, 1, 2, 2, 2, 1], [3, 1, 1, 2, 2, 2, 1]])
compare = [torch.tensor([2, 1]), torch.tensor([2, 3, 1])]
</code></pre>
<p>I want a <code>magic_argsort</code> function give me the following result:</p>
<pre><code class="lang-auto">result = magic_argsort(a, compare)
# result is torch.tensor([[0, 3, 4, 5, 1, 2, 6], [3, 4, 5, 0, 1, 2, 6]])
</code></pre>

username: KFrank
cooked: <p>Hello Usami!</p>
<aside class="quote no-group quote-modified" data-username="usami" data-post="1" data-topic="97235" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/u/b487fb/48.png" class="avatar"> usami:</div>
<blockquote>
<p>How can I specify the sorting by a dynamic comparison order.<br>
…</p>
<p>I want a <code>magic_argsort</code> function give me the following result:</p>
<pre><code class="lang-auto">result = magic_argsort(a, compare)
# result is torch.tensor([[0, 3, 4, 5, 1, 2, 6], [3, 4, 5, 0, 1, 2, 6]])
</code></pre>
</blockquote>
</aside>
<p>I do not believe that pytorch’s <code>sort()</code> supports a custom sort order<br>
or comparator argument.</p>
<p>But you can implement it relatively easily.  The idea will be to look up<br>
your values to be sorted in your comparison-order vectors, and then<br>
sort those indices.</p>
<p>Note that in your example use case, your <code>compare</code> is a python list<br>
of 1D tensors, not a 2D tensor.  Furthermore, the rows differ in length.<br>
Therefore we will have to use a loop, and not just pure pytorch tensor<br>
operations.</p>
<p>Here is a pytorch 0.3.0 version of your <code>magic_argsort()</code> function<br>
applied to your example:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
torch.__version__

a = torch.LongTensor([[2, 1, 1, 2, 2, 2, 1], [3, 1, 1, 2, 2, 2, 1]])
compare = [torch.LongTensor([2, 1]), torch.LongTensor([2, 3, 1])]

# assumes that dataTensor and compareList are equal-length iterables of 1D tensors
# assumes that values in rows of dataTensor do occur in corresponding rows of compareList
def magic_argsort (dataTensor, compareList):
    r = []
    for  d, c  in zip (dataTensor, compareList):
        mask = d.unsqueeze (1) == c   # look up data values; unsqueeze so c can broadcast
        inds = mask.nonzero()[:,1]    # convert to sort-order indices
        r.append (inds.sort()[1])     # 0.3.0 doesn't have separate argsort function
    return torch.stack (r)            # return as torch tensor

magic_argsort (a, compare)
</code></pre>
<p>And here is the output:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'0.3.0b0+591e73e'
&gt;&gt;&gt;
&gt;&gt;&gt; a = torch.LongTensor([[2, 1, 1, 2, 2, 2, 1], [3, 1, 1, 2, 2, 2, 1]])
&gt;&gt;&gt; compare = [torch.LongTensor([2, 1]), torch.LongTensor([2, 3, 1])]
&gt;&gt;&gt;
&gt;&gt;&gt; # assumes that dataTensor and compareList are equal-length iterables of 1D tensors
... # assumes that values in rows of dataTensor do occur in corresponding rows of compareList
... def magic_argsort (dataTensor, compareList):
...     r = []
...     for  d, c  in zip (dataTensor, compareList):
...         mask = d.unsqueeze (1) == c   # look up data values; unsqueeze so c can broadcast
...         inds = mask.nonzero()[:,1]    # convert to sort-order indices
...         r.append (inds.sort()[1])     # 0.3.0 doesn't have separate argsort function
...     return torch.stack (r)            # return as torch tensor
...
&gt;&gt;&gt; magic_argsort (a, compare)

    0     3     4     5     1     2     6
    3     4     5     0     1     2     6
[torch.LongTensor of size 2x7]
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

username: usami
cooked: <p>Thanks K. Frank. I also come up with similar idea without the unsqueeze broadcast trick.<br>
I use this trick in my code now, and speed corresponding code up to 4x times faster.</p>