Tensorflow, tf.multinomial, get the associated probabilities failed

Question

I am trying to using tf.multinomial to sample, and I want to get the associated probability value of the sampled values. Here is my example code,

In [1]: import tensorflow as tf

In [2]: tf.enable_eager_execution()

In [3]: probs = tf.constant([[0.5, 0.2, 0.1, 0.2], [0.6, 0.1, 0.1, 0.1]], dtype=tf.float32)

In [4]: idx = tf.multinomial(probs, 1)

In [5]: idx  # print the indices
Out[5]:
<tf.Tensor: id=43, shape=(2, 1), dtype=int64, numpy=
array([[3],
       [2]], dtype=int64)>

In [6]: probs[tf.range(probs.get_shape()[0], tf.squeeze(idx)]
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-8-56ef51f84ca2> in <module>
----> 1 probs[tf.range(probs.get_shape()[0]), tf.squeeze(idx)]

C:\ProgramData\Anaconda3\lib\site-packages\tensorflow\python\ops\array_ops.py in _slice_helper(tensor, slice_spec, var)
    616       new_axis_mask |= (1 << index)
    617     else:
--> 618       _check_index(s)
    619       begin.append(s)
    620       end.append(s + 1)

C:\ProgramData\Anaconda3\lib\site-packages\tensorflow\python\ops\array_ops.py in _check_index(idx)
    514     # TODO(slebedev): IndexError seems more appropriate here, but it
    515     # will break `_slice_helper` contract.
--> 516     raise TypeError(_SLICE_TYPE_ERROR + ", got {!r}".format(idx))
    517
    518

TypeError: Only integers, slices (`:`), ellipsis (`...`), tf.newaxis (`None`) and scalar tf.int32/tf.int64 tensors are valid indices, got <tf.Tensor: id=7, shape=(2,), dtype=int32, numpy=array([3, 2])>

The expected result I want is [0.2, 0.1] as indicated by idx.

But in Numpy, this method works as answered in https://stackoverflow.com/a/23435869/5046896

How can I fix it?

GoingMyWay · Accepted Answer · 2020-01-30T11:11:06.423

You can try tf.gather_nd, you can try

>>> import tensorflow as tf
>>> tf.enable_eager_execution()
>>> probs = tf.constant([[0.5, 0.2, 0.1, 0.2], [0.6, 0.1, 0.1, 0.1]], dtype=tf.float32)
>>> idx = tf.multinomial(probs, 1)
>>> row_indices = tf.range(probs.get_shape()[0], dtype=tf.int64)
>>> full_indices = tf.stack([row_indices, tf.squeeze(idx)], axis=1)
>>> rs = tf.gather_nd(probs, full_indices)

Or, you can use tf.distributions.Multinomial, the advantage is you do not need to care about the batch_size in the above code. It works under varying batch_size when you set the batch_size=None. Here is a simple example,

multinomail = tf.distributions.Multinomial(
                         total_count=tf.constant(1, dtype=tf.float32),  # sample one for each record in the batch, that is [1, batch_size]
                         probs=probs)
sampled_actions = multinomail.sample()  # sample one action for data in the batch
predicted_actions = tf.argmax(sampled_actions, axis=-1) 
action_probs = sampled_actions * predicted_probs 
action_probs = tf.reduce_sum(action_probs, axis=-1)

I prefer the latter one because it is flexible and elegant.

Tensorflow, tf.multinomial, get the associated probabilities failed

1 Answers1