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