diff --git a/tensorflow_datasets/core/example_parser.py b/tensorflow_datasets/core/example_parser.py index 0fbc9293a6a..0122cb41160 100644 --- a/tensorflow_datasets/core/example_parser.py +++ b/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" diff --git a/tensorflow_datasets/core/example_parser_test.py b/tensorflow_datasets/core/example_parser_test.py index cef10f5a5bb..2e0c92216a9 100644 --- a/tensorflow_datasets/core/example_parser_test.py +++ b/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) diff --git a/tensorflow_datasets/core/features/sequence_feature_test.py b/tensorflow_datasets/core/features/sequence_feature_test.py index ceb0a9798c9..95fc2db6407 100644 --- a/tensorflow_datasets/core/features/sequence_feature_test.py +++ b/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]], ], ),