code: def build_dataset(
      dataset_fn,
      batch_size=None,
      shuffle_buffer_size=256,
      seed=None,
      strategy=None,
      **dataset_kwargs,
  ):
      """Dataset builder that takes care of strategy, batching and shuffling.
  
      Args:
        dataset_fn: function; A function that loads the dataset.
        batch_size: int; Size of the batch.
        shuffle_buffer_size: int; Size of the buffer for used for shuffling.
        seed: int; Random seed used for shuffling.
        strategy: TF strategy that handles the distribution policy.
        **dataset_kwargs: dict; Arguments passed to TFDS.
  
      Returns:
        Dataset.
      """
  
      def _dataset_fn(input_context=None):
          """Dataset function."""
          replica_batch_size = batch_size
          if input_context:
              replica_batch_size = (
                  input_context.get_per_replica_batch_size(batch_size)
              )
          ds = dataset_fn(**dataset_kwargs)
          split = dataset_kwargs.get("split")
          if split == "train":
              # first repeat then shuffle, then batch
              ds = ds.repeat()
              local_seed = seed  # seed for this machine
              if local_seed is not None and input_context:
                  local_seed += input_context.input_pipeline_id
              ds = ds.shuffle(shuffle_buffer_size, seed=local_seed)
              ds = ds.batch(replica_batch_size, drop_remainder=True)
          else:  # test and validation
              # first batch then repeat
              ds = ds.batch(replica_batch_size, drop_remainder=False)
              ds = ds.repeat()
          options = tf.data.Options()
          options.experimental_optimization.parallel_batch = True
          ds = ds.with_options(options)
          return ds.prefetch(tf.data.experimental.AUTOTUNE)
  
      if strategy:
          ds = strategy.experimental_distribute_datasets_from_function(
              _dataset_fn
          )
      else:
          ds = _dataset_fn()
  
      return ds
  
file_path: scenic/scenic/projects/tasseo/dataset_utils.py
imports: from absl import logging
  import jax
  import numpy
  import tensorflow
label_desc: Use image node_id 3 for calling a global function with example usage: build_dataset(dataset_fn, batch_size, shuffle_buffer_size, seed, strategy) and returns: ds, ds