Mathews24
Mathews24

Reputation: 751

tf.where with multiple conditions

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?

Upvotes: 2

Views: 3413

Answers (1)

Laplace Ricky
Laplace Ricky

Reputation: 1687

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)

Upvotes: 2

Related Questions