def tfds_inputs(
    dataset_name,
    preprocess_fun,
    batch_size,
    eval_batch_size=None,
    data_dir=None,
    train_split=tfds.Split.TRAIN,
    eval_split=tfds.Split.VALIDATION,
    shuffle_buffer_size=1024,
    batch_shuffle_size=128,
):
  """Tensorflow Datasets input pipeline, with pure-python preprocessing."""
  if eval_batch_size is None:
    eval_batch_size = batch_size
  return Inputs(
      train_stream=functools.partial(
          _tfds_stream,
          dataset_name=dataset_name,
          split=train_split,
          batch_size=batch_size,
          data_dir=data_dir,
          shuffle_files=True,
          shuffle_buffer_size=shuffle_buffer_size,
          batch_shuffle_size=batch_shuffle_size,
          preprocess_fun=preprocess_fun,
      ),
      eval_stream=functools.partial(
          _tfds_stream,
          dataset_name=dataset_name,
          split=eval_split,
          batch_size=eval_batch_size,
          data_dir=data_dir,
          shuffle_files=False,
          shuffle_buffer_size=None,
          batch_shuffle_size=None,
          preprocess_fun=preprocess_fun,
      ),
  )
Tensorflow Datasets input pipeline, with pure-python preprocessing.
trax/supervised/pretrain_finetune.py