Select a subset of the tensor using the Boolean tensor in the tensor flow

Question

I have two rank-2 tensors arr1and arr2forms mon n. arr2Boolean tensor ; exactly one entry in each of its lines True. I want to extract a new rank-1 tensor of arr3length m, where the ith element arr3is equal to the record from the ith row arr1corresponding to where ith row arr2is equal True.

In numpyI can do it as follows:

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

Is it possible to do something like this in tensorflow? I know I can use eval()my tensors and then use functions numpy, but that seems inefficient.

This question suggests using tf.gatherand tf.select, but this does not concern the folding of the dimension of the output, like my question.

Answer 1

Perhaps you can 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)))

Must give: [2, 3]