tf.where 有多个条件

【问题标题】：tf.where with multiple conditionstf.where 有多个条件
【发布时间】：2021-07-30 17:51:35
【问题描述】：

类似于this question here，给定一个数组dists，可以使用np.where 和& 在多个条件下过滤掉元素，例如 dists[np.where((dists >= r) & (dists**2. <= 100))]

有没有等效的方法在tf.where 中输入多个条件？

【问题讨论】：

标签： python numpy tensorflow

【解决方案1】：

tf.math.logical_and 可用于在 tensorflow 中进行元素和操作。要通过索引从张量中获取指定元素，请使用tf.gather_nd。

示例代码：

@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))

预期输出：

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)

【讨论】：