Repository: tensorflow
Issue: tf.data.experimental.dense_to_ragged_batch fails with inputs from generator with unspecified shape in TF 2.3
Labels: stat:awaiting tensorflower,type:bug,comp:data,TF 2.11

**System information**
- Have I written custom code (as opposed to using a stock example script provided in TensorFlow): Both
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Windows + Linux (Colab)
- Mobile device (e.g. iPhone 8, Pixel 2, Samsung Galaxy) if the issue happens on mobile device: N/A
- TensorFlow installed from (source or binary): binary (PyPI)
- TensorFlow version (use command below): 2.1, 2.2, 2.3
- Python version: 3.6
- Bazel version (if compiling from source): N/A
- GCC/Compiler version (if compiling from source): N/A
- CUDA/cuDNN version: 10.1/7.6
- GPU model and memory: Titan RTX/P100

**Describe the expected behavior**
In TF 2.1, 2.2, and 2.3 batching variable length elements work fine when generated from tensor slices:
```python
ds = tf.data.Dataset.from_tensor_slices(tf.range(4))

# Generate variable length elements via map.
# First batch will have length 1. Subsequent batches will have length 2.
def f(x):
  if x == 0:
    return tf.ones([1,])
  else:
    return tf.ones([2,])
ds = ds.map(f)

# Inspect individual elements.
print("Unbatched shapes:")
for batch in ds:
  print(batch.shape)
print()

# Batch into ragged tensors.
ds = ds.apply(tf.data.experimental.dense_to_ragged_batch(batch_size=2))

# Inspect batched elements.
print("Batched shapes:")
for batch in ds:
  print(batch.to_tensor().shape)
```
Outputs:
```
Unbatched shapes:
(1,)
(2,)
(2,)
(2,)

Batched shapes:
(2, 2)
(2, 2)
```

Now, in TF 2.1 and 2.2, this also works when the dataset consumes elements from a generator:
```python
# Generate elements via a generator.
# First batch will have length 1. Subsequent batches will have length 2.
def gen():
  for i in range(4):
    if i == 0:
      yield tf.ones((1,))
    else:
      yield tf.ones((2,))
ds = tf.data.Dataset.from_generator(gen, output_types=tf.float32)

# Inspect individual elements.
print("Unbatched shapes:")
for batch in ds:
  print(batch.shape)
print()

# Batch into ragged tensors.
ds = ds.apply(tf.data.experimental.dense_to_ragged_batch(batch_size=2))

# Inspect batched elements.
print("Batched shapes:")
for batch in ds:
  print(batch.to_tensor().shape)
```

Outputs:
```
Unbatched shapes:
(1,)
(2,)
(2,)
(2,)

Batched shapes:
(2, 2)
(2, 2)
```
As expected, we get identical outputs both before and after batching.


**Describe the current behavior**
In TF 2.3, the generator version results in an error:
```
[...]
InvalidArgumentError: Cannot batch tensors with different shapes in component 0. First element had shape [1] and element 1 had shape [2]. [Op:IteratorGetNext]

During handling of the above exception, another exception occurred:

InvalidArgumentError                      Traceback (most recent call last)
<ipython-input-4-313ce2c9fab5> in <module>()
     16 
     17 print("Batched shapes:")
---> 18 for batch in ds:
     19   print(batch.to_tensor().shape)
[...]
InvalidArgumentError: Cannot batch tensors with different shapes in component 0. First element had shape [1] and element 1 had shape [2].
```

The release notes for 2.3 mention:

> - tf.data.experimental.dense_to_ragged_batch works correctly with tuples.
> - tf.data.experimental.dense_to_ragged_batch to output variable ragged rank.

Presumably this issue is related to these changes.

Here's the relevant implementation for the actual batching in TF 2.2:
https://github.com/tensorflow/tensorflow/blob/2b96f3662bd776e277f86997659e61046b56c315/tensorflow/python/data/experimental/ops/batching.py#L371-L426

And in TF 2.3:
https://github.com/tensorflow/tensorflow/blob/b36436b087bd8e8701ef51718179037cccdfc26e/tensorflow/python/data/experimental/ops/batching.py#L380-L452

As suggested by the changes above, the behavior prior to 2.3 is achieved again when the output shape is specified, even if unknown:
```python
# Generate elements via a generator.
# First batch will have length 1. Subsequent batches will have length 2.
def gen():
  for i in range(4):
    if i == 0:
      yield tf.ones((1,))
    else:
      yield tf.ones((2,))

# Creating the generator explicitly specifying the unknown shape.
ds = tf.data.Dataset.from_generator(
    gen,
    output_types=tf.float32,
    output_shapes=tf.TensorShape([None])
)

# Inspect individual elements.
print("Unbatched shapes:")
for batch in ds:
  print(batch.shape)
print()

# Batch into ragged tensors.
ds = ds.apply(tf.data.experimental.dense_to_ragged_batch(batch_size=2))

# Inspect batched elements.
print("Batched shapes:")
for batch in ds:
  print(batch.to_tensor().shape)
```

Outputs:
```
Unbatched shapes:
(1,)
(2,)
(2,)
(2,)

Batched shapes:
(2, 2)
(2, 2)
```

It's definitely less convenient to have to specify the output shapes in the generator, requiring some refactoring when updating to 2.3 -- maybe the shapes could just default to unknown when batching if unspecified in the generator?

I appreciate that this is an experimental function, which is noted in the [tf.data.experimental](https://www.tensorflow.org/api_docs/python/tf/data/experimental) docs:

> Note that the tf.data.experimental API is not subject to the same backwards compatibility guarantees as tf.data, but **we will provide deprecation advice in advance of removing existing functionality**.

If this is intended behavior, then perhaps it could be documented somewhere as it does remove the "functionality" of being able to ragged batch elements from a generator without specified output shape :)


**Standalone code to reproduce the issue**
- [Colab: TF 2.1](https://colab.research.google.com/drive/1lm9aOqtGob7PD_LMflmCONuXrpA-M1Kg?usp=sharing) (Works)
- [Colab: TF 2.2](https://colab.research.google.com/drive/1clcicwg22-DVYirFlR_IAlbWaAn-NfUf?usp=sharing) (Works)
- [Colab: TF 2.3](https://colab.research.google.com/drive/1NK5SWwqUaMsisDBjsRe8E-gpc3kcLI5S?usp=sharing) (Fails)
- [Colab: TF 2.3 with `output_shape` specified](https://colab.research.google.com/drive/1WKYYNm96xHIMfdBn_nNwAaN_BmurVuM1?usp=sharing) (Works)


**Other info / logs**
Full traceback of the error in TF 2.3:
```
---------------------------------------------------------------------------
InvalidArgumentError                      Traceback (most recent call last)
/usr/local/lib/python3.6/dist-packages/tensorflow/python/eager/context.py in execution_mode(mode)
   2101       ctx.executor = executor_new
-> 2102       yield
   2103     finally:

/usr/local/lib/python3.6/dist-packages/tensorflow/python/data/ops/iterator_ops.py in _next_internal(self)
    757             output_types=self._flat_output_types,
--> 758             output_shapes=self._flat_output_shapes)
    759 

/usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/gen_dataset_ops.py in iterator_get_next(iterator, output_types, output_shapes, name)
   2609     except _core._NotOkStatusException as e:
-> 2610       _ops.raise_from_not_ok_status(e, name)
   2611     except _core._FallbackException:

/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/ops.py in raise_from_not_ok_status(e, name)
   6842   # pylint: disable=protected-access
-> 6843   six.raise_from(core._status_to_exception(e.code, message), None)
   6844   # pylint: enable=protected-access

/usr/local/lib/python3.6/dist-packages/six.py in raise_from(value, from_value)

InvalidArgumentError: Cannot batch tensors with different shapes in component 0. First element had shape [1] and element 1 had shape [2]. [Op:IteratorGetNext]

During handling of the above exception, another exception occurred:

InvalidArgumentError                      Traceback (most recent call last)
<ipython-input-4-313ce2c9fab5> in <module>()
     16 
     17 print("Batched shapes:")
---> 18 for batch in ds:
     19   print(batch.to_tensor().shape)

/usr/local/lib/python3.6/dist-packages/tensorflow/python/data/ops/iterator_ops.py in __next__(self)
    734 
    735   def __next__(self):  # For Python 3 compatibility
--> 736     return self.next()
    737 
    738   def _next_internal(self):

/usr/local/lib/python3.6/dist-packages/tensorflow/python/data/ops/iterator_ops.py in next(self)
    770   def next(self):
    771     try:
--> 772       return self._next_internal()
    773     except errors.OutOfRangeError:
    774       raise StopIteration

/usr/local/lib/python3.6/dist-packages/tensorflow/python/data/ops/iterator_ops.py in _next_internal(self)
    762         return self._element_spec._from_compatible_tensor_list(ret)  # pylint: disable=protected-access
    763       except AttributeError:
--> 764         return structure.from_compatible_tensor_list(self._element_spec, ret)
    765 
    766   @property

/usr/lib/python3.6/contextlib.py in __exit__(self, type, value, traceback)
     97                 value = type()
     98             try:
---> 99                 self.gen.throw(type, value, traceback)
    100             except StopIteration as exc:
    101                 # Suppress StopIteration *unless* it's the same exception that

/usr/local/lib/python3.6/dist-packages/tensorflow/python/eager/context.py in execution_mode(mode)
   2103     finally:
   2104       ctx.executor = executor_old
-> 2105       executor_new.wait()
   2106 
   2107 

/usr/local/lib/python3.6/dist-packages/tensorflow/python/eager/executor.py in wait(self)
     65   def wait(self):
     66     """Waits for ops dispatched in this executor to finish."""
---> 67     pywrap_tfe.TFE_ExecutorWaitForAllPendingNodes(self._handle)
     68 
     69   def clear_error(self):

InvalidArgumentError: Cannot batch tensors with different shapes in component 0. First element had shape [1] and element 1 had shape [2].
```