title: Multidimensional random draw without replacement with 'predrawn' samples in pytorch
body: <p>I have an <code>(N, I)</code> tensor of <code>N</code> rows with <code>I</code> indices beween <code>0</code> and <code>Z</code>, e.g.,
  <code>N=5, I=3, Z=100</code>:</p>
  
  <pre><code>foo = tensor([[83,  5, 85],
                [ 7, 60, 66],
                [89, 25, 63],
                [58, 67, 47],
                [12, 46, 40]], device='cuda:0')
  </code></pre>
  
  <p>Now I want to efficiently add <code>X</code> random additional <em>new</em> indices (i.e., not yet included in the tensor!) between <code>0</code> and <code>Z</code> to the tensor, e.g.:</p>
  
  <pre><code>foo_new = tensor([[83,  5, 85,  9, 43, 53, 42],
                    [ 7, 60, 66, 85, 64, 22,  1],
                    [89, 25, 63, 38, 24,  4, 75],
                    [58, 67, 47, 83, 43, 29, 55],
                    [12, 46, 40, 74, 21, 11, 52]], device='cuda:0')
  </code></pre>
  
  <p>The tensor would in the end have in each row <code>I+X</code> unique indices between <code>0</code> and <code>Z</code>, where <code>I</code> indices are the ones from the initial tensor, and <code>X</code> indices are uniform randomly drawn without replacement from the remaining indices <code>{0...Z}\{I(n)}</code>, where <code>{I(n)}</code> are the inidices of the nth row.</p>
  
  <p>So it's like a multidimensional random draw without replacement from indices <code>0</code> to <code>Z</code>, where the first <code>I</code> draws (in each row) are enforced to result in the indices given by the initial tensor.</p>
  
  <p>How would I do this efficiently, especially with potentially large <code>Z</code>?</p>
  
  <p>What I tried so far (which was quite slow):</p>
  
  <pre><code>device = torch.cuda.current_device()
  notinfoo = torch.ones((N,I), device=device).byte()
  N_row = torch.arange(N, device=device).unsqueeze(dim=-1)
  notinfoo[N_row, foo] = 0
  foo_new = torch.stack([torch.cat((f, torch.arange(Z, device=device)[nf][torch.randperm(Z-I, device=device)][:X])) for f,nf in zip(foo,notinfoo)])
  </code></pre>