Efficiently selecting a random element from a vector

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?

This code should work:

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]]

I’m not sure if it’s more efficient than just calling torch.randint(0, a.size(0), (1,)) in a while loop to skip the -1 entries. It might depend on the ratio between valid and invalid entries.

1 Like

Hi, ptrblck

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]]

In that case, this should work:

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)]

EDIT: Sorry, I didn’t realized you would like to sample from each line.
Will update the code in a minute.

EDIT2: I’m not sure, if we can avoid using a loop in this case:

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)

Hi, thank you for a great code snippet.

I found an error case where your code returns a Runtime Error.
When there is only one element in valid_idx and that element is at the 0th position in a it returns the following error:

RuntimeError: invalid argument 2: invalid multinomial distribution (sum of probabilities <= 0) at /Users/distiller/project/conda/conda-bld/pytorch_1556653464916/work/aten/src/TH/generic/THTensorRandom.cpp:343

My reproducible code is as follows:

a  = torch.Tensor([-1,4,6,3])
indices = (a == -1).nonzero().view(-1)
choice = torch.multinomial(indices.float(), 1)

Your second suggestion works fine with this case.

I met a similar problem and just figured it out. You can achieve that with no loops:

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)]

Edit: This requires generating many random numbers though, so it might be less efficient if the width of a is very large.