Infinite loop when using reduce() inside tf.data.Dataset.map()

[Elysaveth]

System information

• Have I written custom code: yes
• OS Platform and Distribution: Linux Ubuntu 20.04
• TensorFlow installed from: Binary
• TensorFlow version: v2.4.0-49-g85c8b2a817f 2.4.1
• Python version: 3.8.5
• CUDA/cuDNN version: None
• GPU model and memory: None

Describe the current behavior
I'm trying to create a 3D image dataset using my previous dataset of sorted 2D images. To do so, I use map() method of tf.data.Dataset along with a reduce() function that takes some 2D images from the same dataset (or a copy, from what I tried it doesn't matter) and adds them together. Nevertheless, when I use functools.reduce() or tf.foldl() functions inside of a tf.data.Dataset.map(), even in a simpler case, it enters an infinite loop.
For the example, I've simplified the case to a dummy dataset of 2 elements with shape (2, 3). And I'm using reduce() instead of tf.fold() to avoid using a generator in the input of the function.
The result is also an infinite loop.

Describe the expected behavior
In my understanding, I should get a new dataset with one element of shape (1, 2, 3). But if this transformation is not possible, I think it should raise an error.

Standalone code to reproduce the issue
import tensorflow as tf
from functools import reduce

ds = tf.data.Dataset.from_tensor_slices([[[1,2,3],[4,5,6]],[[7,8,9], [10,11,12]]])
new_ds = ds.map(lambda _: reduce(lambda a, b: tf.concat([tf.expand_dims(a[0], axis=0),tf.expand_dims(b[0], axis=0)], 0), ds))

[Saduf2019]

I am able to replicate this on tf 2.3,2.4 and nightly, please find the gist here

[sushreebarsa]

Was able to replicate the issue in TF 2.11. Please find the gist here..Thanks !