Randomly select items from two equally sized tensors

BramVanroy: <p>Assume that we have two equally sized tensors of size <code>batch_size * 1</code>. For each index in the batch dimension we want to choose randomly between the two tensors. My solution was to create an <code>indices</code> tensor that contains random <code>0</code> or <code>1</code> indices of size <code>batch_size</code> and use those to <code>index_select</code> from the concatenation of the two tensors. However, to do so I had the “view” that <code>cat</code> tensor and the solution ended up to be quite “ugly”:</p>
<pre><code class="lang-python">import torch

bs = 8
a = torch.zeros(bs, 1)
print("a size", a.size())
b = torch.ones(bs, 1)

c = torch.cat([a, b], dim=-1)
print(c)
print("c size", c.size())

# create bs number of random 0 and 1's
indices = torch.randint(0, 2, [bs])
print("idxs size", indices.size())
print("idxs", indices)

# use `indices` to slice the `cat`ted tensor
d = c.view(1, -1).index_select(-1, indices).view(-1, 1)
print("d size", d.size())
print(d)
</code></pre>
<p>I am wondering whether there is a prettier and, more importantly, more efficient solution.</p>

LeviViana: <p>You can do it this way:</p>
<pre><code class="lang-python">import torch

bs = 8
a = torch.zeros(bs, 1)
print("a size", a.size())
b = torch.ones(bs, 1)

idx = torch.randint(2 * bs, (bs,))

d = torch.cat([a, b])[idx] # [bs, 1]
</code></pre>
<p>d is a [bs, 1 ] tensor containing a random sample of both batchs, with possible sample duplicates, which is OK for most statistical applications, it is just a sort of bootstraping.</p>

hhaoao: <p>According to this link, this may be the fastest way.<br>
</p><aside class="onebox githubissue">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/issues/16897" target="_blank" rel="nofollow noopener">github.com/pytorch/pytorch</a>
  </header>
  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue">
	  <svg width="60" height="60" class="github-icon" viewBox="0 0 14 16" aria-hidden="true"><path d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/issues/16897" target="_blank" rel="nofollow noopener">Implement `numpy.random.choice` equivalent</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2019-02-08" data-time="15:33:11" data-timezone="UTC">03:33PM - 08 Feb 19 UTC</span>
      </div>


      <div class="user">
        <a href="https://github.com/LeviViana" target="_blank" rel="nofollow noopener">
          <img alt="LeviViana" src="https://avatars1.githubusercontent.com/u/17881241?v=4" class="onebox-avatar-inline" width="20" height="20">
          LeviViana
        </a>
      </div>
    </div>
  </div>
</div>

<div class="github-row">
  <p class="github-content">🚀 Feature
Implement numpy.random.choice equivalent.
Motivation
In some cases, it is useful to get random samples from a torch Tensor efficiently. For instance, during...</p>
</div>

<div class="labels">
    <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">feature</span>
    <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">module: numpy</span>
    <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">topic: bootcamp</span>
    <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">triaged</span>
</div>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

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

bs = 8
a = torch.zeros(bs, 1)
b = torch.ones(bs, 1)
c = torch.cat([a, b], dim=-1)
choices_flat = c.view(-1)
# index = torch.randint(choices_flat.numel(), (bs,))
# or if replace = False
index = torch.randperm(choices_flat.numel())[:bs]
select = choices_flat[index]

print(select)
</code></pre>