[Python] TensorFlow `next()` hangs in infinite loop after `take(1)`

When using TensorFlow's tf.data.Dataset.take(1) to get the first element of a dataset and then attempting to retrieve this element using iter() and next(), the next() call enters an infinite loop if the pandas library is imported before tensorflow.

This behavior has been observed in both VS Code Jupyter Notebook and standard Jupyter Notebook environments, suggesting an interaction at the Python runtime level.

Code to reproduce:

import pandas as pd
import tensorflow as tf

# Assume your original dataset looks something like this
def generator():
for i in range(5):
yield (tf.random.uniform(shape=(2, 7, 11)), tf.random.uniform(shape=(2, 1, 1)))

original_dataset = tf.data.Dataset.from_generator(
generator,
output_signature=(
tf.TensorSpec(shape=(None, 7, 11), dtype=tf.float32),
tf.TensorSpec(shape=(None, 1, 1), dtype=tf.float32)
)
)

# take the first element
first_element_ds = original_dataset.take(1)
iterator = iter(first_element_ds)

# get the first element with next() will run into an infinite loop
first_element_tuple = next(iterator)

tensor1, tensor2 = first_element_tuple
print("Tensor 1:", tensor1.numpy().shape)
print("Tensor 2:", tensor2.numpy().shape)


Environment:

• Pandas Version: 2.2.2
• TensorFlow Version: 2.17.0
• Python Version: 3.11.9
• NumPy Version: 1.26.4
• Jupyter Core Version: 5.7.2
• Operating System: macOS 13.7.4
• Virtual Environment: conda 24.11.0, cPython 3.11.5

Question:

Why does the the execution of next() runs into an infinite loop?

Continue reading...