Add tf.RaggedTensor support for tfds.as_numpy

PiperOrigin-RevId: 278750600

Author: Conchylicultor

tensorflow_datasets/core/dataset_utils.py

@@ -159,7 +159,7 @@ def _build_ds_from_instruction(instruction, ds_from_file_fn):
 def _eager_dataset_iterator(dataset):
   for item in dataset:
     flat = tf.nest.flatten(item)
-    flat = [el.numpy() for el in flat]
+    flat = [t if isinstance(t, tf.RaggedTensor) else t.numpy() for t in flat]
     yield tf.nest.pack_sequence_as(item, flat)
 
 
@@ -184,6 +184,13 @@ def as_numpy(dataset, graph=None):
   `as_numpy` converts a possibly nested structure of `tf.data.Dataset`s
   and `tf.Tensor`s to iterables of NumPy arrays and NumPy arrays, respectively.
 
+  Note that because TensorFlow has support for ragged tensors and NumPy has
+  no equivalent representation,
+  [`tf.RaggedTensor`s](https://www.tensorflow.org/api_docs/python/tf/RaggedTensor)
+  are left as-is for the user to deal with them (e.g. using `to_list()`).
+  In TF 1 (i.e. graph mode), `tf.RaggedTensor`s are returned as
+  `tf.ragged.RaggedTensorValue`s.
+
   Args:
     dataset: a possibly nested structure of `tf.data.Dataset`s and/or
       `tf.Tensor`s.
@@ -204,7 +211,9 @@ def as_numpy(dataset, graph=None):
   for ds_el in flat_ds:
     types = [type(el) for el in flat_ds]
     types = tf.nest.pack_sequence_as(nested_ds, types)
-    if not (isinstance(ds_el, tf.Tensor) or tf_compat.is_dataset(ds_el)):
+    if not (
+        isinstance(ds_el, (tf.Tensor, tf.RaggedTensor)) or
+        tf_compat.is_dataset(ds_el)):
       raise ValueError("Arguments to as_numpy must be tf.Tensors or "
                        "tf.data.Datasets. Got: %s" % types)
 
@@ -213,6 +222,8 @@ def as_numpy(dataset, graph=None):
     for ds_el in flat_ds:
       if isinstance(ds_el, tf.Tensor):
         np_el = ds_el.numpy()
+      elif isinstance(ds_el, tf.RaggedTensor):
+        np_el = ds_el
       elif tf_compat.is_dataset(ds_el):
         np_el = _eager_dataset_iterator(ds_el)
       else:

tensorflow_datasets/core/dataset_utils_test.py

@@ -135,6 +135,50 @@ def test_tensors_match(self):
     # and target may not match
     self.assertAllEqual(ds["a"], ds["b"])
 
+  @testing.run_in_graph_and_eager_modes()
+  def test_ragged_tensors(self):
+    rt = tf.ragged.constant([
+        [1, 2, 3],
+        [],
+        [4, 5],
+    ])
+    rt = dataset_utils.as_numpy(rt)
+
+    if not tf.executing_eagerly():
+      # Output of `sess.run(rt)` is a `RaggedTensorValue` object
+      self.assertIsInstance(rt, tf.compat.v1.ragged.RaggedTensorValue)
+    else:
+      self.assertIsInstance(rt, tf.RaggedTensor)
+
+    self.assertAllEqual(rt, tf.ragged.constant([
+        [1, 2, 3],
+        [],
+        [4, 5],
+    ]))
+
+  @testing.run_in_graph_and_eager_modes()
+  def test_ragged_tensors_ds(self):
+    def _gen_ragged_tensors():
+      # Yield the (flat_values, rowids)
+      yield ([0, 1, 2, 3], [0, 0, 0, 2])  # ex0
+      yield ([], [])  # ex1
+      yield ([4, 5, 6], [0, 1, 1])  # ex2
+    ds = tf.data.Dataset.from_generator(
+        _gen_ragged_tensors,
+        output_types=(tf.int64, tf.int64),
+        output_shapes=((None,), (None,))
+    )
+    ds = ds.map(tf.RaggedTensor.from_value_rowids)
+
+    rt0, rt1, rt2 = list(dataset_utils.as_numpy(ds))
+    self.assertAllEqual(rt0, [
+        [0, 1, 2],
+        [],
+        [3,],
+    ])
+    self.assertAllEqual(rt1, [])
+    self.assertAllEqual(rt2, [[4], [5, 6]])
+
 
 class DatasetOffsetTest(testing.TestCase):
   """Test that the offset functions are working properly."""