python - 带有 bool 张量的 tensorflow 索引

Question

在numpy中，有两个相同形状的数组x和y，可以像这样y[x>1]。您如何在 tensorflow 中获得相同的结果？ y[tf.greater(x, 1)] 不起作用，tf.slice 也不支持这样的东西。现在有没有办法使用 bool 张量进行索引，还是目前不受支持？

Answer 1

试试:

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)

见 tf.boolean_mask

编辑:另一种(更好的？)方法:

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]