CODEWITHSUNDEEP
tensorflow
Bob
Bob

Reputation: 571

Sample from a tensor in Tensorflow along an axis

I have a matrix L of shape (2,5,2). The values along the last axis form a probability distribution. I want to sample another matrix S of shape (2, 5) where each entry is one of the following integers: 0, 1. For example,

L = [[[0.1, 0.9],[0.2, 0.8],[0.3, 0.7],[0.5, 0.5],[0.6, 0.4]],
     [[0.5, 0.5],[0.9, 0.1],[0.7, 0.3],[0.9, 0.1],[0.1, 0.9]]]

One of the samples could be,

S = [[1, 1, 1, 0, 1],
     [1, 1, 1, 0, 1]]

The distributions are binomial in the above example. However, in general, the last dimension of L can be any positive integer, so the distributions can be multinomial.

The samples need to be generated efficiently within Tensorflow computation graph. I know how to do this using numpy using the functions apply_along_axis and numpy.random.multinomial.

Upvotes: 3

Views: 4826

Answers (2)

sunder
sunder

Reputation: 1

Be cautious when using tf.multinomial(). The inputs to the function should be logits and not probability distributions. However, in your example, the last axis is a probability distribution.

Upvotes: 0

Olivier Moindrot
Olivier Moindrot

Reputation: 28218

You can use tf.multinomial() here.

You will first need to reshape your input tensor to shape [-1, N] (where N is the last dimension of L):

# L has shape [2, 5, 2]
L = tf.constant([[[0.1, 0.9],[0.2, 0.8],[0.3, 0.7],[0.5, 0.5],[0.6, 0.4]],
                 [[0.5, 0.5],[0.9, 0.1],[0.7, 0.3],[0.9, 0.1],[0.1, 0.9]]])

dims = L.get_shape().as_list()
N = dims[-1]  # here N = 2

logits = tf.reshape(L, [-1, N])  # shape [10, 2]

Now we can apply the function tf.multinomial() to logits:

samples = tf.multinomial(logits, 1)
# We reshape to match the initial shape minus the last dimension
res = tf.reshape(samples, dims[:-1])

Upvotes: 6

Related Questions