Efficiently selecting a random element from a vector

Naruto-Sasuke: <pre><code class="lang-python">import torch
# Assuming that I have a tensor a.
a  = torch.Tensor([-1,4,6,3,-1,-1])
# Then I need to select a random element(not -1) from a.
# How to do that efficiently?
</code></pre>

ptrblck: <p>This code should work:</p>
<pre><code class="lang-python">a  = torch.Tensor([-1,4,6,3,-1,-1])
valid_idx = (a!=-1).nonzero().view(-1)
choice = torch.multinomial(valid_idx.float(), 1)
a[valid_idx[choice]]
</code></pre>
<p>I’m not sure if it’s more efficient than just calling <code>torch.randint(0, a.size(0), (1,))</code> in a while loop to skip the <code>-1</code> entries. It might depend on the ratio between valid and invalid entries.</p>

Naruto-Sasuke: <p>Hi, ptrblck</p>
<pre><code class="lang-auto">import torch
# Assuming that I have a tensor a.
a = torch.Tensor([[3,4, -1, -1,5,6],
                  [-1,5,-1,1,-1,-1],
                  [-1,4,4,2,3,0]])
# Then I need to select a random element(not -1) from a for each line
# How to do that efficiently?
# Eg.
# a_sample = [[4],[1],[3]]
</code></pre>

ptrblck: <p>In that case, this should work:</p>
<pre><code class="lang-python">a = torch.Tensor([[3,4, -1, -1,5,6],
                  [-1,5,-1,1,-1,-1],
                  [-1,4,4,2,3,0]])

valid_idx = (a!=-1).nonzero()
choice = torch.multinomial(torch.arange(valid_idx.size(0)).float(), 1)
a[valid_idx[choice].squeeze().chunk(2)]
</code></pre>
<p>EDIT: Sorry, I didn’t realized you would like to sample from each line.<br>
Will update the code in a minute.</p>
<p>EDIT2: I’m not sure, if we can avoid using a loop in this case:</p>
<pre><code class="lang-python">valid_idx = (a!=-1).nonzero()
unique_rows = valid_idx[:, 0].unique()
valid_row_idx = [valid_idx[valid_idx[:, 0] == u] for u in unique_rows]

ret = []
for v in valid_row_idx:
    choice = torch.multinomial(torch.arange(v.size(0)).float(), 1)
    ret.append(a[v[choice].squeeze().chunk(2)])
ret = torch.stack(ret)
</code></pre>

jsideas: <p>Hi, thank you for a great code snippet.</p>
<p>I found an error case where your code returns a Runtime Error.<br>
When there is only one element in <code>valid_idx</code> and that element is at the 0th position in <code>a</code> it returns the following error:</p>
<pre><code class="lang-auto">RuntimeError: invalid argument 2: invalid multinomial distribution (sum of probabilities &lt;= 0) at /Users/distiller/project/conda/conda-bld/pytorch_1556653464916/work/aten/src/TH/generic/THTensorRandom.cpp:343
</code></pre>
<p>My reproducible code is as follows:</p>
<pre><code class="lang-auto">a  = torch.Tensor([-1,4,6,3])
indices = (a == -1).nonzero().view(-1)
choice = torch.multinomial(indices.float(), 1)
</code></pre>
<p>Your second suggestion works fine with this case.</p>

guannanzhao: <p>I met a similar problem and just figured it out. You can achieve that with no loops:</p>
<pre><code class="lang-auto">a = torch.Tensor([[3,4, -1, -1,5,6],
                  [-1,5,-1,1,-1,-1],
                  [-1,4,4,2,3,0]])

b = torch.rand_like(a)
b[a == -1] = -1
a[torch.arange(a.shape[0]), b.argmax(dim=1)]
</code></pre>
<p>Edit: This requires generating <em>many random numbers</em> though, so it might be less efficient if the width of <code>a</code> is very large.</p>