Fix parsing of nested sequences with fixed lengths

tensorflow_datasets/core/example_parser.py

@@ -220,6 +220,20 @@ def _feature_to_numpy(
     return value_array.item()
   feature_name = _key_to_feature_name(key)
   row_lengths = _extract_row_lengths(features, feature_name)
+  if tensor_info.sequence_rank > 1:
+    # Every nested Sequence has stored row lengths, even with a fixed length.
+    # Keep the trailing tensor dimensions together while rebuilding sequences.
+    values: Union[np.ndarray, list[Any]] = value_array.reshape(
+        (sum(row_lengths[-1]),) + shape[tensor_info.sequence_rank :]
+    )
+    for lengths in reversed(row_lengths):
+      rows = []
+      offset = 0
+      for length in lengths:
+        rows.append(values[offset : offset + length])
+        offset += length
+      values = rows
+    return values
   return reshape_ragged_tensor(value_array, row_lengths, shape)  # pyrefly: ignore[bad-argument-type]
 
 
@@ -380,23 +394,7 @@ def _to_tf_example_spec(tensor_info: feature_lib.TensorInfo):
   # Convert the shape
 
   # Select the feature proto type in function of the unknown shape
-  if all(s is not None for s in tensor_info.shape):
-    return tf.io.FixedLenFeature(  # All shaped defined
-        shape=tensor_info.shape,
-        dtype=dtype,
-        default_value=tensor_info.default_value,
-    )
-  elif tensor_info.shape.count(None) == 1:
-    # Extract the defined shape (without the None dimension)
-    # The original shape is restored in `_deserialize_single_field`
-    shape = tuple(dim for dim in tensor_info.shape if dim is not None)
-    return tf.io.FixedLenSequenceFeature(  # First shape undefined
-        shape=shape,
-        dtype=dtype,
-        allow_missing=True,
-        default_value=tensor_info.default_value,
-    )
-  elif tensor_info.sequence_rank > 1:  # RaggedTensor
+  if tensor_info.sequence_rank > 1:  # RaggedTensor
     # Decoding here should match encoding from `_add_ragged_fields` in
     # `example_serializer.py`
     tf_specs = {  # pylint: disable=g-complex-comprehension
@@ -414,6 +412,22 @@ def _to_tf_example_spec(tensor_info: feature_lib.TensorInfo):
         default_value=tensor_info.default_value,
     )
     return tf_specs
+  elif all(s is not None for s in tensor_info.shape):
+    return tf.io.FixedLenFeature(  # All shaped defined
+        shape=tensor_info.shape,
+        dtype=dtype,
+        default_value=tensor_info.default_value,
+    )
+  elif tensor_info.shape.count(None) == 1:
+    # Extract the defined shape (without the None dimension)
+    # The original shape is restored in `_deserialize_single_field`
+    shape = tuple(dim for dim in tensor_info.shape if dim is not None)
+    return tf.io.FixedLenSequenceFeature(  # First shape undefined
+        shape=shape,
+        dtype=dtype,
+        allow_missing=True,
+        default_value=tensor_info.default_value,
+    )
   else:
     raise NotImplementedError(
         "Multiple unknown dimension not supported.\n"

tensorflow_datasets/core/example_parser_test.py

@@ -18,7 +18,9 @@
 import numpy as np
 import pytest
 from tensorflow_datasets.core import example_parser
+from tensorflow_datasets.core import example_serializer
 from tensorflow_datasets.core import features as features_lib
+from tensorflow_datasets.core.utils.lazy_imports_utils import tensorflow as tf
 from tensorflow_datasets.proto import tf_example_pb2
 from tensorflow_datasets.proto import tf_feature_pb2
 
@@ -84,7 +86,7 @@ def test_example_parser_np():
       'array_of_unknown_shape': np.asarray([4, 5, 6], dtype=np.uint16),
       'ragged_tensor': [
           [[0, 1], [2, 3]],
-          [],
+          np.empty((0, 2), dtype=np.int32),
           [[4, 5]],
       ],
       'ragged_tensor_without_length': [
@@ -175,3 +177,81 @@ def test_full_row_lengths():
           (None, 3, None),
       )
   ) == [[3, 0, 2, 2, 2, 1], [3, 3]]
+
+
+@pytest.mark.parametrize(
+    'lengths',
+    [
+        (outer, inner)
+        for outer in (None, 0, 2)
+        for inner in (None, 0, 2)
+    ] + [
+        (outer, middle, inner)
+        for outer in (None, 0, 2)
+        for middle in (None, 0, 2)
+        for inner in (None, 0, 2)
+    ],
+)
+@pytest.mark.parametrize('tensor_shape', [(), (3,), (2, 3)])
+@pytest.mark.parametrize('parser_type', ['numpy', 'eager', 'graph'])
+def test_nested_sequence_round_trip(lengths, tensor_shape, parser_type):
+  feature = features_lib.Tensor(shape=tensor_shape, dtype=np.int16)
+  for length in reversed(lengths):
+    feature = features_lib.Sequence(feature, length=length)
+  features = features_lib.FeaturesDict({'value': feature})
+
+  def make_value(level, index=0):
+    if level == len(lengths):
+      return np.full(tensor_shape, index, dtype=np.int16)
+    length = lengths[level]
+    if length is None:
+      length = 2 if level == 0 else index % 3
+    return [make_value(level + 1, i) for i in range(length)]
+
+  value = make_value(0)
+  specs = features.get_serialized_info()
+  serialized = example_serializer.ExampleSerializer(specs).serialize_example(
+      features.encode_example({'value': value})
+  )
+  if parser_type == 'numpy':
+    actual = example_parser.ExampleParserNp(specs).parse_example(serialized)[
+        'value'
+    ]
+
+    def assert_rows(actual_rows, expected_rows, level):
+      assert len(actual_rows) == len(expected_rows)
+      if level == len(lengths) - 1:
+        assert isinstance(actual_rows, np.ndarray)
+        assert actual_rows.dtype == np.int16
+        assert actual_rows.shape == (len(expected_rows),) + tensor_shape
+        np.testing.assert_array_equal(
+            actual_rows,
+            np.asarray(expected_rows, dtype=np.int16).reshape(
+                (len(expected_rows),) + tensor_shape
+            ),
+        )
+      else:
+        for actual_row, expected_row in zip(actual_rows, expected_rows):
+          assert_rows(actual_row, expected_row, level + 1)
+
+    assert_rows(actual, value, 0)
+  else:
+    parser = example_parser.ExampleParser(specs)
+    parse = parser.parse_example
+    if parser_type == 'graph':
+      parse = tf.function(parse)
+    actual = parse(serialized)['value']
+    assert isinstance(actual, tf.RaggedTensor)
+    assert actual.dtype == tf.int16
+    assert actual.ragged_rank == len(lengths) - 1
+    assert actual.shape.as_list() == (
+        [len(value)] + [None] * (len(lengths) - 1) + list(tensor_shape)
+    )
+    assert actual.flat_values.shape.as_list()[1:] == list(tensor_shape)
+
+    def as_list(item):
+      if isinstance(item, list):
+        return [as_list(child) for child in item]
+      return item.tolist()
+
+    assert actual.to_list() == as_list(value)

tensorflow_datasets/core/features/sequence_feature_test.py

@@ -325,7 +325,7 @@ def test_2lvl_sequences(self):
                 ),
                 expected_np=[
                     [[0, 1], [2, 3]],
-                    [],
+                    np.empty((0, 2), dtype=np.int32),
                     [[4, 5]],
                 ],
             ),
@@ -339,7 +339,9 @@ def test_2lvl_sequences(self):
             testing.FeatureExpectationItem(
                 value=[[], [], []],
                 expected=[[], [], []],
-                expected_np=[[], [], []],
+                expected_np=[
+                    np.empty((0, 2), dtype=np.int32) for _ in range(3)
+                ],
             ),
             # List of empty np.array
             testing.FeatureExpectationItem(
@@ -352,8 +354,8 @@ def test_2lvl_sequences(self):
                     [],
                 ],
                 expected_np=[
-                    [],
-                    [],
+                    np.empty((0, 2), dtype=np.int32),
+                    np.empty((0, 2), dtype=np.int32),
                 ],
             ),
             testing.FeatureExpectationItem(
@@ -368,8 +370,8 @@ def test_2lvl_sequences(self):
                     [[1, 1], [1, 1], [1, 1]],
                 ],
                 expected_np=[
-                    [],
-                    [],
+                    np.empty((0, 2), dtype=np.int32),
+                    np.empty((0, 2), dtype=np.int32),
                     [[1, 1], [1, 1], [1, 1]],
                 ],
             ),