Sample from two torch.tensors based on certain weights or probabilities

#1

Hello,
I’m curious if its possible to do something like np.random.choice([a, b], p=[p_a, p_b])
using PyTorch where a and b are 1-d tensors of length L and p_a and p_b are probabilities used to sample elements from either tensor a or b into the resultant 1-d tensor.

I think torch.multinomial([p_a, p_b], L) might be useful here - it returns a 1-d tensor length L of 0s and 1s based on the probabilities I give it but I’m drawing a blank of how to utilize this tensor to sample from my tensors a and b. Appreciate any help!

#2

I’m not sure to completely understand the use case, but this code snippet should work, if you would like to sample either the element from a or b based on the probabilities:

L = 10
a = torch.arange(L)
b = torch.arange(L, 2*L)

p_a = 0.2
probs = torch.empty(1, L).uniform_()
idx = (probs >= p_a).to(torch.long)
c = torch.stack((a, b))
result = torch.gather(c, 0, idx)

print(c)
> tensor([[ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9],
          [10, 11, 12, 13, 14, 15, 16, 17, 18, 19]])
print(probs)
> tensor([[0.8268, 0.0361, 0.2744, 0.4381, 0.7244, 0.8509, 0.5047, 0.0549, 0.2767,           0.3203]])
print(idx)
> tensor([[1, 0, 1, 1, 1, 1, 1, 0, 1, 1]])
print(result)
> tensor([[10,  1, 12, 13, 14, 15, 16,  7, 18, 19]])
3 Likes
#3

Thank you - this is great! I ended up using the gather function(which I was previously unaware of) to sample from my a and b tensors after stacking them in dim=0 with the indices I got from torch.multinomial()

1 Like