Batched permutation

javivilladamigo: <p>Hi there,</p>
<p>I am trying to permute a tensor that is shaped [batch_shape, 4, 4] along its third dimension. The idea is that I want to get a random permutation of (0,1,2,3) indices, i.e. (0,1,3,2), (0,2,1,3), …, (3,1,2,0); and permute each of the elements along the batch with it.</p>
<p>For example, first element may be indexed such that</p>
<pre><code class="lang-auto">permuted_tensor[0] = tensor_to_permute[0, :, (3,1,2,0)]
</code></pre>
<p>and second elment, such that:</p>
<pre><code class="lang-auto">permuted_tensor[1] = tensor_to_permute[1, :, (1,2,0,3)]
</code></pre>
<p>Now this is one of those things that would be much easier with a for loop:</p>
<pre><code class="lang-auto">for i in range(tensor_to_permute.shape[0]):
     permuted_tensor[i] = tensor_to_permute[i, :, torch.randperm(4)]
</code></pre>
<p>but I am very reluctant to use it, since the computational time increases notably. Any advice on how to do it?</p>
<p>Thanks!</p>

KFrank: <p>Hi Javier!</p>
<aside class="quote no-group" data-username="javivilladamigo" data-post="1" data-topic="184339" 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/javivilladamigo/48/62291_2.png" class="avatar"> javivilladamigo:</div>
<blockquote>
<pre><code class="lang-auto">for i in range(tensor_to_permute.shape[0]):
     permuted_tensor[i] = tensor_to_permute[i, :, torch.randperm(4)]
</code></pre>
</blockquote>
</aside>
<p>The problem is that pytorch (still) does not offer a batch version of<br>
<code>randperm()</code>.  The standard work-around is to compute <code>argsort()</code><br>
for a batch of random vectors.</p>
<p>Thus:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; _ = torch.manual_seed (2023)
&gt;&gt;&gt;
&gt;&gt;&gt; tensor_to_permute = torch.randint (100, (3, 4, 4))   # some test data
&gt;&gt;&gt; perms = torch.rand (3, 4).argsort (dim = 1)          # a way to get a batch of permutations
&gt;&gt;&gt;
&gt;&gt;&gt; permuted_tensor = torch.empty_like (tensor_to_permute)
&gt;&gt;&gt; for i in range(tensor_to_permute.shape[0]):
...     permuted_tensor[i] = tensor_to_permute[i, :, perms[i]]
...
&gt;&gt;&gt; permuted_tensorB = tensor_to_permute.gather (2, perms.unsqueeze (1).expand (-1, 4, -1))
&gt;&gt;&gt;
&gt;&gt;&gt; torch.equal (permuted_tensor, permuted_tensorB)
True
</code></pre>
<p>(You could use pytorch tensor indexing (“advanced indexing”) instead of<br>
<code>gather()</code>, if you prefer that syntax.)</p>
<p>This is mildly inefficient in that the cost of sorting a vector of length <code>n</code> is<br>
<code>O (n log (n))</code> while the cost of generating a random permutation of<br>
length <code>n</code> is <code>O (n)</code>.  In your example, <code>n = 4</code>, so that <code>log (n)</code> factor is<br>
irrelevant, but as <code>n</code> becomes large, it could matter.  However, the slowdown<br>
due to the python loop will be worse than the <code>log (n)</code> unless <code>batch_shape</code><br>
is small and <code>n</code> is quite larger.</p>
<p>Best.</p>
<p>K. Frank</p>

Craig_Hicks1: <p>There is a problem with using argsort to generate a batch of permutations.  That problem is explained above the source code for <code>torch.randperm</code>.  I quote it here:</p>
<blockquote>
<p>// [Algorithm of randperm]<br>
//<br>
// randperm is implemented by sorting an arange tensor of size n with randomly<br>
// generated keys. When random keys are different from each other, all different<br>
// permutations have the same probability.<br>
//<br>
// However, there is a pitfall here:<br>
// For better performance, these N random keys are generated independently,<br>
// and there is no effort to make sure they are different at the time of generation.<br>
// When two keys are identical, stable sorting algorithms will not permute these two keys.<br>
// As a result, (0, 1) will appear more often than (1, 0).<br>
//<br>
// To overcome this pitfall we first carefully choose the number of bits in these keys,<br>
// so that the probability of having duplicate keys is under a threshold. Let q be the<br>
// threshold probability for having non-duplicate keys, then it can be proved that[1]<br>
// the number of bits required is: ceil(log2(n - (6 n^2 + 1) / (12 log(q))))<br>
//<br>
// Then after sort, we lauch a separate kernel that additionally shuffles any islands<br>
// of values whose keys matched. The algorithm of this kernel is as follows:<br>
// Each thread reads its key and the keys of its neighbors to tell if it’s part of an island.<br>
// For each island, the first thread in the island sees a key match at index i+1 but not index i-1.<br>
// This thread considers itself the “island leader”. The island leader then reads more indices to<br>
// the right to figure out how big the island is. Most likely, the island will be very small,<br>
// just a few values. The island leader then rolls that many RNG, uses them to additionally<br>
// shuffle values within the island using serial Fisher-Yates, and writes them out.<br>
//<br>
// Reference<br>
// [1] <a href="https://osf.io/af2hy/" class="inline-onebox" rel="noopener nofollow ugc">OSF</a></p>
</blockquote>
<p><a href="https://github.com/pytorch/pytorch/blob/c4ae45104f87e96dcd6d5ba47cb4e356fc2c63f3/aten/src/ATen/native/cuda/Randperm.cu#L23" rel="noopener nofollow ugc">Here is a link to the above comment.</a></p>
<p>So you can see from the comment that <code>randperm</code> addresses the problem so that the permutation is truly random.  (Well, maybe, because the same small chance of collisions exists again in the solution unless that is checked for).  If you don’t use <code>randperm</code>, there is higher chance of adjacent index pairs in ascending order than descending order.  E.g.  <code>...12...</code> will occur more often than <code>...21...</code>.</p>

Chenwei_Cui: <pre><code class="lang-auto">import torch

def batched_randperm(batch_size, seq_len, device):
    out = torch.randperm(batch_size * seq_len, dtype=torch.int64, device=device)
    out = out.view(batch_size, seq_len)
    out = out.argsort(dim=1)
    return out
</code></pre>
<p>This approach is parallelizable and collision-free.</p>
<p>In a uniformly random permutation, the relative order of elements in any subsequence is itself uniformly random.</p>