Better Tensor singleton repr

PiperOrigin-RevId: 297170026

Author: Conchylicultor

tensorflow_datasets/core/features/feature.py

@@ -516,7 +516,7 @@ class Tensor(FeatureConnector):
   @api_utils.disallow_positional_args
   def __init__(self, shape, dtype):
     """Construct a Tensor feature."""
-    self._shape = shape
+    self._shape = tuple(shape)
     self._dtype = dtype
 
   def get_tensor_info(self):
@@ -544,3 +544,25 @@ def encode_example(self, example_data):
           example_data.dtype, np_dtype))
     utils.assert_shape_match(example_data.shape, self._shape)
     return example_data
+
+
+def get_inner_feature_repr(feature):
+  """Utils which returns the object which should get printed in __repr__.
+
+  This is used in container features (Sequence, FeatureDict) to print scalar
+  Tensor in a less verbose way `Sequence(tf.int32)` rather than
+  `Sequence(Tensor(shape=(), dtype=tf.in32))`.
+
+  Args:
+    feature: The feature to dispaly
+
+  Returns:
+    Either the feature or it's inner value.
+  """
+  # We only print `tf.int32` rather than `Tensor(shape=(), dtype=tf.int32)`
+  # * For the base `Tensor` class (and not subclass).
+  # * When shape is scalar (explicit check to avoid trigger when `shape=None`).
+  if type(feature) == Tensor and feature.shape == ():  # pylint: disable=unidiomatic-typecheck,g-explicit-bool-comparison
+    return repr(feature.dtype)
+  else:
+    return repr(feature)

tensorflow_datasets/core/features/features_dict.py

@@ -141,7 +141,8 @@ def __repr__(self):
     lines = ['{}({{'.format(type(self).__name__)]
     # Add indentation
     for key, feature in sorted(list(self._feature_dict.items())):
-      all_sub_lines = '\'{}\': {},'.format(key, feature)
+      feature_repr = feature_lib.get_inner_feature_repr(feature)
+      all_sub_lines = '\'{}\': {},'.format(key, feature_repr)
       lines.extend('    ' + l for l in all_sub_lines.split('\n'))
     lines.append('})')
     return '\n'.join(lines)

tensorflow_datasets/core/features/features_test.py

@@ -22,6 +22,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import textwrap
 import numpy as np
 import tensorflow.compat.v2 as tf
 from tensorflow_datasets import testing
@@ -206,7 +207,16 @@ def test_feature__repr__(self):
         'label': features_lib.Sequence(label),
     })
 
-    self.assertEqual(repr(feature_dict), FEATURE_STR)
+    self.assertEqual(
+        repr(feature_dict),
+        textwrap.dedent("""\
+        FeaturesDict({
+            'label': Sequence(ClassLabel(shape=(), dtype=tf.int64, num_classes=2)),
+            'metadata': Sequence({
+                'frame': Image(shape=(32, 32, 3), dtype=tf.uint8),
+            }),
+        })"""),
+    )
 
   def test_feature_save_load_metadata_slashes(self):
     with testing.tmp_dir() as data_dir:
@@ -218,14 +228,6 @@ def test_feature_save_load_metadata_slashes(self):
       fd.load_metadata(data_dir)
 
 
-FEATURE_STR = """FeaturesDict({
-    'label': Sequence(ClassLabel(shape=(), dtype=tf.int64, num_classes=2)),
-    'metadata': Sequence({
-        'frame': Image(shape=(32, 32, 3), dtype=tf.uint8),
-    }),
-})"""
-
-
 class FeatureTensorTest(testing.FeatureExpectationsTestCase):
 
   def test_shape_static(self):
@@ -395,6 +397,39 @@ def test_string(self):
         ],
     )
 
+  def test_repr_tensor(self):
+
+    # Top level Tensor is printed expanded
+    self.assertEqual(
+        repr(features_lib.Tensor(shape=(), dtype=tf.int32)),
+        'Tensor(shape=(), dtype=tf.int32)',
+    )
+
+    # Sequences colapse tensor repr
+    self.assertEqual(
+        repr(features_lib.Sequence(tf.int32)),
+        'Sequence(tf.int32)',
+    )
+
+    class ChildTensor(features_lib.Tensor):
+      pass
+
+    self.assertEqual(
+        repr(features_lib.FeaturesDict({
+            'colapsed': features_lib.Tensor(shape=(), dtype=tf.int32),
+            # Tensor with defined shape are printed expanded
+            'noncolapsed': features_lib.Tensor(shape=(1,), dtype=tf.int32),
+            # Tensor inherited are expanded
+            'child': ChildTensor(shape=(), dtype=tf.int32),
+        })),
+        textwrap.dedent("""\
+        FeaturesDict({
+            'child': ChildTensor(shape=(), dtype=tf.int32),
+            'colapsed': tf.int32,
+            'noncolapsed': Tensor(shape=(1,), dtype=tf.int32),
+        })"""),
+    )
+
 
 if __name__ == '__main__':
   testing.test_main()

tensorflow_datasets/core/features/sequence_feature.py

@@ -201,7 +201,7 @@ def __setstate__(self, state):
 
   def __repr__(self):
     """Display the feature."""
-    inner_feature_repr = repr(self._feature)
+    inner_feature_repr = feature_lib.get_inner_feature_repr(self._feature)
     if inner_feature_repr.startswith('FeaturesDict('):
       # Minor formatting cleaning: 'Sequence(FeaturesDict({' => 'Sequence({'
       inner_feature_repr = inner_feature_repr[len('FeaturesDict('):-len(')')]