python - numpy random choice in Tensorflow

Question

Is there an equivalent function to numpy random choice in Tensorflow. In numpy we can get an item randomly from the given list with its weights.

 np.random.choice([1,2,3,5], 1, p=[0.1, 0, 0.3, 0.6, 0])

This code will select an item from the given list with p weights.

Answer 1

No, but you can achieve the same result using tf.multinomial:

elems = tf.convert_to_tensor([1,2,3,5])
samples = tf.multinomial(tf.log([[1, 0, 0.3, 0.6]]), 1) # note log-prob
elems[tf.cast(samples[0][0], tf.int32)].eval()
Out: 1
elems[tf.cast(samples[0][0], tf.int32)].eval()
Out: 5

The [0][0] part is here, as multinomial expects a row of unnormalized log-probabilities for each element of the batch and also has another dimension for the number of samples.