def _tfds_stream(n_devices,
                 dataset_name,
                 split,
                 batch_size,
                 data_dir,
                 shuffle_files,
                 shuffle_buffer_size,
                 batch_shuffle_size,
                 preprocess_fun,
                 repeat=True):
  """Streams batches of examples from tfds, with pure-python preprocessing."""
  # TODO(piotrekp1): delete if switched to data_streams
  if batch_size % n_devices != 0:
    raise ValueError(f'Batch size ({batch_size}) not divisible'
                     ' by number of devices ({n_devices})')
  ds = tfds.load(
      name=dataset_name,
      split=split,
      data_dir=data_dir,
      shuffle_files=shuffle_files)
  if repeat:
    ds = ds.repeat()
  if shuffle_buffer_size is not None:
    ds = ds.shuffle(shuffle_buffer_size)
  ds = ds.batch(batch_size)
  if batch_shuffle_size is not None:
    ds = ds.shuffle(batch_shuffle_size)

  for batch in tfds.as_numpy(ds):
    if preprocess_fun is not None:
      yield preprocess_fun(batch)
    else:
      yield batch
Streams batches of examples from tfds, with pure-python preprocessing.
trax/supervised/pretrain_finetune.py