dataset.map does not return / runs indefinitely in graph mode?

Viewed 41

The following code creates a tensorflow Dataset with 3 elements and runs a trivial function on its elements using map. Importantly, the results are then converted to a Python list.

import tensorflow as tf
dataset = tf.data.Dataset.from_tensor_slices(tf.constant(range(3)))
res_map = dataset.map(
    lambda x: x*2
)

# @tf.function # hangs in graph mode
def outer():
    return tf.convert_to_tensor(list(res_map))
res = outer()
print(res)

runs fine in eager mode and returns

tf.Tensor([0 2 4], shape=(3,), dtype=int32)

as expected.

However, when I uncomment the decorator to run outer in graph mode, the line res = outer() does not return (it runs for at least 10 minutes on my machine).

Is there a fundamental problem with my code or is this a (possibly known) bug / shortcoming of tensorflow? In the first case, what should I do differently? I am using tensorflow 2.2.0.

0 Answers
Related