code: 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
  
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 1 for calling a global function with example usage: _tfds_stream(n_devices, dataset_name, split, batch_size, data_dir, shuffle_files, shuffle_buffer_size, batch_shuffle_size, preprocess_fun, repeat) without return types