tf.boolean_mask

import tensorflow as tf

tensor = [0, 1, 2, 3]  # 1-D example
mask = np.array([True, False, True, False])
c=tf.boolean_mask(tensor, mask)
print(c)

print('------------------------------')

tensor = [[1, 2], [3, 4], [5, 6]] # 2-D example
mask = np.array([True, False, True])
c=tf.boolean_mask(tensor, mask)
print(c)

tf.Tensor([0 2], shape=(2,), dtype=int32)
------------------------------
tf.Tensor(
[[1 2]
 [5 6]], shape=(2, 2), dtype=int32)