Loop over a tensor in Tensorflow

user_11970084 0
Viewed 600

I want to loop over a Tensorflow tensor, my code is something like this:

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)

How can I do it in Tensorflow under graph mode? I got following error:

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

I'm running it with TF 2.4.1 but I am doing it under Beam context (for TFX purpose) which means the operations are done under graph operation.

Thanks!

tfx apache-beam tensorflow python
1 Answers

It's solved this way:

    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)
user_119700840
Related