Shuffling every slice of a tensor using different permutations

wasabi: <p>I have a tensor X of shape (a, b, c) and a matrix of permutation (not a permutation matrix) P of shape (a,b), where each row of P is an output of torch.randperm(). I want to shuffle X in the following way:</p>
<pre><code class="lang-auto">for i in range(a):
    Y[i] = X[i][P[i]]
return Y
</code></pre>
<p>what’s the best way to achieve this? Thanks.</p>

KFrank: <p>Hi Wasabi!</p>
<aside class="quote no-group" data-username="wasabi" data-post="1" data-topic="127880" 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/wasabi/48/13151_2.png" class="avatar"> wasabi:</div>
<blockquote>
<pre><code class="lang-auto">for i in range(a):
    Y[i] = X[i][P[i]]
return Y
</code></pre>
</blockquote>
</aside>
<p>Use <code>gather()</code>:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'1.9.0'
&gt;&gt;&gt; _ = torch.manual_seed (2021)
&gt;&gt;&gt; a = 2
&gt;&gt;&gt; b = 3
&gt;&gt;&gt; c = 5
&gt;&gt;&gt; X = torch.arange (a * b * c).reshape (a, b, c)
&gt;&gt;&gt; X
tensor([[[ 0,  1,  2,  3,  4],
         [ 5,  6,  7,  8,  9],
         [10, 11, 12, 13, 14]],

        [[15, 16, 17, 18, 19],
         [20, 21, 22, 23, 24],
         [25, 26, 27, 28, 29]]])
&gt;&gt;&gt; P = torch.randn (a, b).argsort (dim = 1)
&gt;&gt;&gt; P
tensor([[2, 1, 0],
        [1, 0, 2]])
&gt;&gt;&gt; Y = X.gather (1, P.unsqueeze (-1).repeat (1, 1, c))
&gt;&gt;&gt; Y
tensor([[[10, 11, 12, 13, 14],
         [ 5,  6,  7,  8,  9],
         [ 0,  1,  2,  3,  4]],

        [[20, 21, 22, 23, 24],
         [15, 16, 17, 18, 19],
         [25, 26, 27, 28, 29]]])
&gt;&gt;&gt; _ = Y.zero_()
&gt;&gt;&gt; for i in range(a):
...     Y[i] = X[i][P[i]]
...
&gt;&gt;&gt; Y
tensor([[[10, 11, 12, 13, 14],
         [ 5,  6,  7,  8,  9],
         [ 0,  1,  2,  3,  4]],

        [[20, 21, 22, 23, 24],
         [15, 16, 17, 18, 19],
         [25, 26, 27, 28, 29]]])
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

wasabi: <p>Hi KFrank, thanks for your answer. Is there a way that I can use gather() without repeat P along certain axis? Say that my X is a tensor with more axis (a, b, c, d, …), then is there a way to naturally broadcast P?</p>

KFrank: <p>Hi Wasabi!</p>
<aside class="quote no-group" data-username="wasabi" data-post="3" data-topic="127880" 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/wasabi/48/13151_2.png" class="avatar"> wasabi:</div>
<blockquote>
<p>is there a way to naturally broadcast P?</p>
</blockquote>
</aside>
<p>No.  Quoting from the documentation for <a href="https://pytorch.org/docs/1.9.0/generated/torch.gather.html" rel="noopener nofollow ugc">torch.gather()</a>:</p>
<blockquote>
<p>Note that <code>input</code> and <code>index</code> do not broadcast against each other.</p>
</blockquote>
<p>(Also, I’m not aware of any way to do what you want with pytorch<br>
indexing operations other than <code>gather()</code>, although I may have<br>
overlooked something.)</p>
<p>Best.</p>
<p>K. Frank</p>

tom: <p>I’m old enough to remember what we did before broadcasting, and that was to use <code>.expand</code>. Expand will create a view(!) with stride 0 for the expanded dimension, so it is much more efficient that repeat. (Personally, I don’t think I remember when I last saw a use of <code>repeat</code> that should not be <code>expand</code> instead.)</p>
<pre><code class="lang-auto">Y = X.gather(1, P.unsqueeze(-1).expand(-1, -1, c))
</code></pre>
<p>Gather is the most elegant way to do it by far, but there also is the option of explicitly indexing which might be the most immediate transcription of the loop you wrote:</p>
<pre><code class="lang-auto">i0 = torch.arange(a)[:, None, None].expand(-1, b, c)
i1 = P.unsqueeze(-1).expand(-1, -1, c)
i2 = torch.arange(c)[None, None, :].expand(a, b, -1)
print((X[i0, i1, i2] == Y).all())
</code></pre>
<p>(But I don’t know where the third dimension came from, you can leave that out to get closer to the initial message.)</p>
<p>Best regards</p>
<p>Thomas</p>

wasabi: <p>Thanks! I guess it’s not that bad to repeat axis after all.</p>