在张量流中使用布尔张量选择张量的子集

Question

我有两个 rank-2 张量 arr1 和 arr2，形状为 m by n。张量 arr2 是布尔值；在其每一行中恰好有一个条目是True。我想提取长度为m 的新秩1 张量arr3，其中arr3 的ith 条目等于arr1 的ith 行中的条目，对应于i第arr2行等于True。

在numpy，我可以这样做：

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

我可以在tensorflow 中做类似的事情吗？我知道我可以eval() 我的张量然后使用numpy 函数，但这似乎效率低下。

这个question 建议使用tf.gather 和tf.select，但它不像我的问题那样处理折叠输出的维度。

Answer 1

你也许可以使用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)))

应该给出：[2, 3]