code: 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,
          ),
      )
  
file_path: trax/trax/supervised/pretrain_finetune.py
imports: import functools
  import gin
  import numpy
  import tensorflow_datasets
  from trax.data.inputs import Inputs
label_desc: Use image node_id 2 for calling a global function with example usage: tfds_inputs(dataset_name, preprocess_fun, batch_size, eval_batch_size, data_dir, train_split, eval_split, shuffle_buffer_size, batch_shuffle_size) and returns: Inputs