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
Speed up tfds.as_numpy by batching and then iterating over the batches.
trax/fastmath/jax.py