TensorFlow equivalent of np.in1d

Question

I am trying to do:

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

result = [True,False,False,False,True,False]

which is the np.in1d function https://docs.scipy.org/doc/numpy/reference/generated/numpy.in1d.html

Is there a way to implement this in TensorFlow?

Thanks!

Asked By: op10no4

||

Source

Answer 1

You can use tf.equal with broadcasting to form 5x2 matrix where i,j entry has True if a[i]==b[j] and then tf.reduce_any to collapse to bool vector

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)

Answered By: Yaroslav Bulatov

TensorFlow equivalent of np.in1d

Question:

Answers: