在 Tensorflow 中循环一个张量

Question

我想循环一个 Tensorflow 张量，我的代码是这样的：

elements = tf.constant([1,2,3])
x = tf.constant([1.000001, 1.1, 2.1, 2.00004, 3.001])
EPSILON = 0.0001
for elem in elements:
    mask = tf.experimental.numpy.isclose(x, elem, atol=EPSILON, rtol=0)
    x = tf.boolean_mask(x, ~mask)

在Tensorflow中如何在图形模式下做到这一点？我收到以下错误：

    OperatorNotAllowedInGraphError: iterating over `tf.Tensor` is not allowed: AutoGraph did convert this function. This might indicate you are trying to use an unsupported feature.

我正在使用 TF 2.4.1 运行它，但我在 Beam 上下文中执行它（用于 TFX 目的），这意味着操作是在图形操作下完成的。

谢谢！

Answer 1

这样解决了：

    mask = tf.map_fn(fn=lambda t: ~tf.experimental.numpy.isclose(
        x, t, atol=EPSILON, rtol=0
    ), elems=elements)

    mask = tf.reduce_min(mask, axis=0)
    x = tf.boolean_mask(x, mask)