tf.where with multiple conditions

Question

Akin to this question here, given an array dists, one can filter out elements following multiple conditions using np.where and &, e.g. dists[np.where((dists >= r) & (dists**2. <= 100))]

Is there an equivalent way to input more than one condition in tf.where?

score 2 · Accepted Answer · answered Aug 04 '21 at 08:21

tf.math.logical_and can be used to do element-wise and operation in tensorflow. To get specified elements out of a tensor by their indices, use tf.gather_nd.

Example codes:

@tf.function
def eg_func(x,a,b):
  return tf.gather_nd(x,tf.where(tf.math.logical_and(x > a,x < b)))

x=tf.reshape(tf.range(3*3),[3,3])
print(x)
print(eg_func(x,2,7))

Expected outputs:

tf.Tensor(
[[0 1 2]
 [3 4 5]
 [6 7 8]], shape=(3, 3), dtype=int32)
tf.Tensor([3 4 5 6], shape=(4,), dtype=int32)

tf.where with multiple conditions

1 Answers1