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.