How to quickly inverse a permutation by using PyTorch?

111414: <p>I am confused on how to quickly restore an array shuffled by a permutation.</p>
<pre><code class="lang-auto">Example #1:
[x, y, z] shuffled by P: [2, 0, 1], we will obtain [z, x, y],
the corresponding inverse should be P^-1: [1, 2, 0].

Example #2:
[a, b, c, d, e, f] shuffled by P: [5, 2, 0, 1, 4, 3], then we will get [f, c, a, b, e, d],
the corresponding inverse should be P^-1: [2, 3, 1, 5, 4, 0].
</code></pre>
<p>I wrote the following codes based on matrix multiplication (the transpose of permutation matrix is its inverse), but this approach is too slow when I utilize it on my model training. Does there exisits a faster implementation?</p>
<pre><code class="lang-auto">import torch

n = 10
x = torch.Tensor(list(range(n)))
print('Original array', x)

random_perm_indices = torch.randperm(n).long()
perm_matrix = torch.eye(n)[random_perm_indices].t()
x = x[random_perm_indices]
print('Shuffled', x)

restore_indices = torch.Tensor(list(range(n))).view(n, 1)
restore_indices = perm_matrix.mm(restore_indices).view(n).long()
x = x[restore_indices]
print('Restored', x)
</code></pre>

KFrank: <p>Hi Bin!</p>
<aside class="quote no-group" data-username="111414" data-post="1" data-topic="116205" 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/111414/48/31166_2.png" class="avatar"> 111414:</div>
<blockquote>
<pre><code class="lang-auto">Example #1:
[x, y, z] shuffled by P: [2, 0, 1], we will obtain [z, x, y],
the corresponding inverse should be P^-1: [1, 2, 0].

Example #2:
[a, b, c, d, e, f] shuffled by P: [5, 2, 0, 1, 4, 3], then we will get [f, c, a, b, e, d],
the corresponding inverse should be P^-1: [2, 3, 1, 5, 4, 0].
</code></pre>
</blockquote>
</aside>
<p>If you represent your permutation as a vector of permuted integers<br>
(as you do), you may use <code>argsort()</code> to obtain the inverse permutation:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'1.7.1'
&gt;&gt;&gt; p1 = torch.tensor ([2, 0, 1])
&gt;&gt;&gt; torch.argsort (p1)
tensor([1, 2, 0])
&gt;&gt;&gt; p2 = torch.tensor ([5, 2, 0, 1, 4, 3])
&gt;&gt;&gt; torch.argsort (p2)
tensor([2, 3, 1, 5, 4, 0])
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

111414: <p>Oh, that’s what I need. Thank you very much. <img src="https://discuss.pytorch.org/images/emoji/apple/smile.png?v=10" title=":smile:" class="emoji" alt=":smile:"></p>

tom: <p><a class="mention" href="/u/kfrank">@KFrank</a> 's method is a great ad hoc way to inverse a permutation.</p>
<p>If the argument is rather large (say &gt;=10000 elements) and you know it is a permutation (0…9999) then you could also use indexing:</p>
<pre><code class="lang-python">def inverse_permutation(perm):
    inv = torch.empty_like(perm)
    inv[perm] = torch.arange(perm.size(0), device=perm.device)
    return inv
</code></pre>
<p>This will give the same result (again, if you knew the input is a permutation), but on my computer and with 100.000 elements, a quick <code>%timeit</code> benchmark has it faster on the CPU (~4ms vs. ~165µs) and on the GPU (~147µs vs. ~20µs).<br>
In other words, this is something where the O(n log n) complexity of sort vs. the O(n) complexity of index assignment shows for large operands.</p>
<p>The beauty of <a class="mention" href="/u/kfrank">@KFrank</a>’s solution is that you don’t need to write a new function, it instantly generalizes to batches, too. Unless you are in the CPU, have large operands and are in a hurry or in a very tight loop, the speed difference probably doesn’t matter much.</p>
<p>Best regards</p>
<p>Thomas</p>

KFrank: <p>Hi Thomas (and Bin)!</p>
<aside class="quote no-group" data-username="tom" data-post="4" data-topic="116205" 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/tom/48/3162_2.png" class="avatar"> tom:</div>
<blockquote>
<p>this is something where the O(n log n) complexity of sort vs. the O(n) complexity of index assignment shows for large operands.</p>
</blockquote>
</aside>
<p>This is absolutely correct – using <code>argsort()</code> to invert a permutation<br>
has the sub-optimal O(n log n) complexity, and will definitely matter with<br>
large permutations.</p>
<p>Best.</p>
<p>K. Frank</p>

Haydnspass: <p>Just want to add that if the dim indices are negative (which they could be because we are talking python) one needs to account for that.</p>
<p>So my solution would be:</p>
<pre><code class="lang-python">def invert_permutation(*dims: int):
    n = len(dims)
    dims = [d if d &gt;= 0 else d + n for d in dims]
    return torch.argsort(torch.LongTensor(dims)).tolist()

# passes this test
@pytest.mark.parametrize("p_in,p_expct", [
    ([0, 1, 2], [0, 1, 2]),  # identity
    ([0, 2, 1], [0, 2, 1]),  # swap last 2
    ([-3, -2, 1], [0, 1, 2]),  # identity with neg. ix
    ([0, -1, -2], [0, 2, 1]),  # swap last 2
])
def test_inverse_permute(p_in, p_expct):
    assert lazy.tensor.invert_permutation(*p_in) == p_expct


</code></pre>

harpone: <p>Here’s a batch version a la <a class="mention" href="/u/tom">@tom</a> above:</p>
<p>Suppose</p>
<pre><code class="lang-auto">perms = torch.tensor([[3, 2, 1, 0],
                      [1, 2, 3, 0]])
src = torch.arange(4)[None].repeat(2, 1)
perms_inv = torch.empty_like(perms)
perms_inv = torch.scatter(perms_inv, dim=1, index=perms, src=src)
</code></pre>
<p>Output is:</p>
<pre><code class="lang-auto">tensor([[3, 2, 1, 0],
        [3, 0, 1, 2]])
</code></pre>
<p>Requires a <code>.repeat</code> of the arange since gather/scatter don’t support broadcasting, but that’s a minor annoyance.</p>

tom: <p>Thank you for amending the example.<br>
A minor nit: likely <code>.expand</code> (to the right size) has a chance of being a tad faster than repeat.</p>
<p>Best regards</p>
<p>Thomad</p>

harpone: <p>oh yeah, just came back here to point that out! <img src="https://discuss.pytorch.org/images/emoji/apple/sweat_smile.png?v=12" title=":sweat_smile:" class="emoji" alt=":sweat_smile:" loading="lazy" width="20" height="20"> also that expand doesn’t spend extra memory since it returns a view, when repeat returns a copy.</p>