# 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 = [,,]
``````

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.