The dataset_utils.as_numpy doesn't convert RaggedTensors. Fix that but providing custom conversion.

PiperOrigin-RevId: 736137874

Author: The TensorFlow Datasets Authors

tensorflow_datasets/robotics/asimov/asimov.py

@@ -18,11 +18,11 @@
 from tensorflow_datasets.robotics import dataset_importer_builder
 
 ASIMOV_CITATION = """
-@article{Sermanet2025RobotConstitutions,
+@article{sermanet2025asimov,
   author    = {Pierre Sermanet and Anirudha Majumdar and Alex Irpan and Dmitry Kalashnikov and Vikas Sindhwani},
   title     = {Generating Robot Constitutions & Benchmarks for Semantic Safety},
-  journal   = {arXiv preprint arXiv:FIXME},
-  url       = {https://arxiv.org/abs/FIXME},
+  journal   = {arXiv preprint arXiv:2503.08663},
+  url       = {https://arxiv.org/abs/2503.08663},
   year      = {2025},
 }
 """

tensorflow_datasets/robotics/dataset_importer_builder.py

@@ -28,6 +28,7 @@
 from tensorflow_datasets.core.utils import read_config as read_config_lib
 from tensorflow_datasets.core.utils.lazy_imports_utils import tensorflow as tf
 import tensorflow_datasets.public_api as tfds
+import tree
 
 
 
@@ -218,7 +219,7 @@ def _generate_examples(
     read_config = read_config_lib.ReadConfig(add_tfds_id=True)
 
     def converter_fn(example):
-      example_out = dataset_utils.as_numpy(example)
+      example_out = tree.map_structure(to_np, example)
       example_id = example_out['tfds_id'].decode('utf-8')
       del example_out['tfds_id']
 
@@ -235,3 +236,10 @@ def get_ds_builder(self):
     ds_location = self.get_dataset_location()
     ds_builder = tfds.builder_from_directory(ds_location)
     return ds_builder
+
+
+def to_np(tensor):
+  """Convert tensor to numpy."""
+  if isinstance(tensor, tf.Tensor) or isinstance(tensor, tf.RaggedTensor):
+    return tensor.numpy()
+  return tensor