Fix np decoding for Dataset feature.

PiperOrigin-RevId: 647270132

Author: pierrot0

tensorflow_datasets/core/features/dataset_feature.py

@@ -24,7 +24,6 @@
 from tensorflow_datasets.core.data_sources import python
 from tensorflow_datasets.core.features import feature as feature_lib
 from tensorflow_datasets.core.features import sequence_feature
-from tensorflow_datasets.core.features import tensor_feature
 from tensorflow_datasets.core.features import top_level_feature
 from tensorflow_datasets.core.utils import py_utils
 from tensorflow_datasets.core.utils import type_utils
@@ -66,7 +65,7 @@ class Dataset(sequence_feature.Sequence):
 
   ```python
     features=tfds.features.FeatureDict({
-     'agent_id': np.object_,
+      'agent_id': np.object_,
       'episode': tfds.features.Dataset({
         'observation': tfds.features.Image(),
         'reward': tfds.features.Image(),
@@ -176,23 +175,13 @@ def decode_example_np(
     flatten = self.feature._flatten  # pylint: disable=protected-access
     nest = self.feature._nest  # pylint: disable=protected-access
     flat_example = flatten(serialized_example)
-    flat_features = flatten(self.feature)
     num_slices: int | None = None
 
-    # First discover the number of slices in the Dataset. Notably, it's possible
-    # that tensors have to be reshaped. We call slice a record in the Dataset.
+    # Discover the number of slices in the Dataset (ie: the outter dimension).
+    # We call slice a record in the Dataset.
     # We don't use `example` to avoid confusion with the `serialized_example`.
-    for i, feature in enumerate(flat_features):
-      if isinstance(feature, tensor_feature.Tensor) and feature.shape:
-        try:
-          flat_example[i] = flat_example[i].reshape((-1,) + feature.shape)
-        except ValueError as e:
-          raise ValueError(
-              "The length of all elements of one slice should be the same."
-          ) from e
-        feature_num_slices = flat_example[i].shape[0]
-      else:
-        feature_num_slices = len(flat_example[i])
+    for example_feature in flat_example:
+      feature_num_slices = len(example_feature)
       if num_slices is None:
         num_slices = feature_num_slices
       else:

tensorflow_datasets/core/features/dataset_feature_test.py

@@ -492,11 +492,9 @@ def test_getattr(self):
     )
     self.assertEqual(feature.names, ['left', 'right'])
 
-    feature = feature_lib.Dataset(
-        {
-            'label': feature_lib.ClassLabel(names=['left', 'right']),
-        }
-    )
+    feature = feature_lib.Dataset({
+        'label': feature_lib.ClassLabel(names=['left', 'right']),
+    })
     self.assertEqual(feature['label'].names, ['left', 'right'])
 
   def test_metadata(self):
@@ -512,24 +510,52 @@ def test_metadata(self):
 
 class DecodeExampleNpTest(testing.SubTestCase):
 
-  def test_top_level_feature(self):
-    feature = feature_lib.Dataset(
-        {'feature_name': feature_lib.Tensor(dtype=np.uint8, shape=(4, 2))}
-    )
-    example = {'feature_name': np.ones(shape=(24,), dtype=np.int32)}
-    expected = [{'feature_name': np.ones(shape=(4, 2), dtype=np.int32)}] * 3
-    self.assertAllEqualNested(feature.decode_example_np(example), expected)
-
-  def test_tensor_feature(self):
-    feature = feature_lib.Dataset(
-        feature_lib.Tensor(dtype=np.uint8, shape=(4, 2))
-    )
-    example = np.ones(shape=(24,), dtype=np.uint8)
-    expected = [np.ones(shape=(4, 2), dtype=np.int32)] * 3
-    self.assertAllEqualNested(feature.decode_example_np(example), expected)
+  def test_representative_example(self):
+    feature = feature_lib.FeaturesDict({
+        'step_number': feature_lib.Tensor(dtype=np.int32, shape=()),
+        'steps': feature_lib.Dataset({
+            'tensor': feature_lib.Tensor(dtype=np.uint8, shape=(7, 8)),
+            'strings': feature_lib.Tensor(dtype=np.str_, shape=(3,)),
+            'bool': feature_lib.Tensor(dtype=np.bool_, shape=()),
+            'obj': feature_lib.FeaturesDict({
+                'a': feature_lib.Tensor(
+                    dtype=np.float32,
+                    shape=(5,),
+                    encoding=feature_lib.Encoding.ZLIB,
+                ),
+                'b': feature_lib.Tensor(dtype=np.int32, shape=(6,)),
+            }),
+            'reward': feature_lib.Tensor(dtype=np.float32, shape=()),
+        }),
+        'timestamp': feature_lib.Tensor(dtype=np.int64, shape=()),
+    })
+    subdataset_size = 42
+    example = {
+        'step_number': 7,
+        'steps': [
+            {
+                'tensor': np.ones(shape=(7, 8), dtype=np.uint8),
+                'strings': ['foo', 'bar', 'baz'],
+                'bool': True,
+                'obj': {
+                    'a': np.zeros(shape=(5,), dtype=np.float32),
+                    'b': np.zeros(shape=(6,), dtype=np.int32),
+                },
+                'reward': np.float32(42.42),
+            }
+            for _ in range(subdataset_size)
+        ],
+        'timestamp': 1234567890,
+    }
+    encoded_example = feature.encode_example(example)
+    decoded_encoded_example = feature.decode_example_np(encoded_example)
+    self.assertAllEqualNested(decoded_encoded_example, example)
 
   def test_nested_dict(self):
-    feature = feature_lib.Dataset({'a': {'b': np.int32}, 'b': np.str_})
+    feature = feature_lib.Dataset({
+        'a': {'b': np.int32},
+        'b': np.str_,
+    })
     example = {'a': {'b': [1, 2, 3]}, 'b': ['a', 'b', 'c']}
     expected = [
         {'a': {'b': 1}, 'b': 'a'},