code: def _dataset_as_numpy(ds, batch_size=None):
      """Speed up tfds.as_numpy by batching and then iterating over the batches."""
      batch_size = batch_size or 1
      try:  # Check that dense_to_ragged_batch exists.
          if (
              batch_size < 2
          ):  # Fall back to default if no batching requested.
              raise AttributeError
          ds_batch = ds.apply(
              tf.data.experimental.dense_to_ragged_batch(batch_size)
          )
          for example in tfds.as_numpy(ds_batch):
              flat_example = tnp.tree_flatten(example)
              np_flat_example = [_to_numpy(x) for x in flat_example]
              for single_example_flat in zip(*np_flat_example):
                  single_example, _ = tnp.tree_unflatten(
                      single_example_flat, example
                  )
                  yield single_example
      except AttributeError:
          # In TF 1.X there is not dense_to_ragged_batch: fallback.
          for example in tfds.as_numpy(ds):
              yield example
  
imports: import functools
  import jax
  from jax import lax
  from jax import random
  import jax.numpy
  import jax.scipy.special
  import numpy
  import tensorflow
  import tensorflow_datasets
  from trax.fastmath import numpy
  from trax.shapes import signature
label_desc: Use image node_id 10 for calling a global function with example usage: _dataset_as_numpy(ds, batch_size) without return types