TensorFlow相当于np.in1d

Question

我正在尝试：

a = [1,2,3,4,5,6] b = [1,5]

result = [True，False，False，False，True，False]

这是np.in1d函数https://docs.scipy.org/doc/numpy/reference/generated/numpy.in1d.html

有没有办法在TensorFlow中实现这个？

谢谢！

Answer 1

您可以使用tf.equal进行广播以形成5x2矩阵，其中i,j条目True a[i]==b[j]，然后tf.reduce_any折叠为a = [1,2,3,4,5,6] b = [1,5] a0 = tf.expand_dims(a, 1) b0 = tf.expand_dims(b, 0) result = sess.run(tf.reduce_any(tf.equal(a0, b0), 1)) assert result == np.in1d(a, b) 布尔矢量

$CWD