Select subset of tensor using boolean tensor in tensorflow

Question

I have two rank-2 tensors arr1 and arr2 of shape m by n. The tensor arr2 is boolean; precisely one entry in each of its rows is True. I want to extract a new rank-1 tensor arr3 of length m, where the ith entry of arr3 is equal to the entry from the ith row of arr1 corresponding to where the ith row of arr2 is equal to True.

In numpy, I can do this as follows:

arr1 = np.array([[1,2],
                 [3,4]])
arr2 = np.array([[0,1],
                 [1,0]], dtype="bool")
arr3 = arr1[arr2]

Can I do something similar in tensorflow? I know that I could eval() my tensors and then use numpy functions, but that seems inefficient.

This question suggests the use of tf.gather and tf.select, but it does not deal with collapsing the dimensionality of the output like my question does.

score 2 · Accepted Answer · answered Feb 18 '17 at 07:48

You could maybe use tf.boolean_mask?

from __future__ import print_function

import tensorflow as tf

with tf.Session() as sess:
    arr1 = tf.constant([[1,2],
                        [3,4]])
    arr2 = tf.constant([[False, True],
                        [True, False]])
    print(sess.run(tf.boolean_mask(arr1, arr2)))

Should give: [2, 3]

Select subset of tensor using boolean tensor in tensorflow

1 Answers1