python - Tensorflow indexing with boolean tensor

Question

In numpy, with two arrays of the same shape, x and y, it is possible to do slices like this y[x > 1]. How do you achieve the same result in tensorflow? y[tf.greater(x, 1)] doesn't work and tf.slice doesn't support anything like this either. Is there a way to index with a boolean tensor right now or is that currently unsupported?

Answer 1

Try:

ones = tf.ones_like(x) # create a tensor all ones
mask = tf.greater(x, ones) # boolean tensor, mask[i] = True iff x[i] > 1
slice_y_greater_than_one = tf.boolean_mask(y, mask)

See tf.boolean_mask

EDIT: another (better ?) way to do it:

import tensorflow as tf

x = tf.constant([1, 2, 0, 4])
y = tf.Variable([1, 2, 0, 4])
mask = x > 1
slice_y_greater_than_one = tf.boolean_mask(y, mask)

with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    print (sess.run(slice_y_greater_than_one)) # [2 4]