code: def get_dataset(
      *,
      batch_size,
      eval_batch_size,
      num_shards,
      dtype_str="float32",
      shuffle_seed=0,
      rng=None,
      prefetch_buffer_size=2,
      dataset_configs=None,
      dataset_service_address: Optional[str] = None,
  ):
      """Returns generators for the pileup window train, validation, and test set.
  
      Args:
        batch_size: int; Determines the train batch size.
        eval_batch_size: int; Determines the evaluation batch size.
        num_shards: int;  Number of shards --> batch shape: [num_shards, bs, ...].
        dtype_str: Data type of the image (e.g. 'float32').
        shuffle_seed: int; Seed for shuffling the training data.
        rng: JAX rng key, which can be used for augmentation, shuffling, etc.
        prefetch_buffer_size: int; Buffer size for the prefetch.
        dataset_configs: dict; Dataset specific configurations.
        dataset_service_address: If set, will distribute the training dataset using
          the given tf.data service at the given address.
  
      Returns:
        A dataset_utils.Dataset() which includes a train_iter, a valid_iter,
          a test_iter, and a dict of meta_data.
      """
      del rng
  
      def build_pileup_window_dataset(split):
          """dataset_fn called by data.build_dataset(**kwargs)."""
  
          if split == "train":
              if dataset_configs.train_path:
                  path = dataset_configs.train_path
              else:
                  # Add path to your training data here:
                  path = ""
          elif split == "valid":
              if dataset_configs.eval_path:
                  path = dataset_configs.eval_path
              else:
                  # Add path to your test data here:
                  path = ""
          elif split == "test":
              if dataset_configs.test_path:
                  path = dataset_configs.test_path
              else:
                  # Add path to your test data here:
                  path = ""
          if not path:
              raise ValueError(
                  "No path provide. Please modify the path variable to "
                  "hardcode the %s path." % split
              )
          dataset_files = tf.io.matching_files(path)
  
          # Creating a Dataset that includes only 1/num_shards of data so the data is
          # splitted between different hosts.
          num_hosts, host_id = jax.process_count(), jax.process_index()
  
          if len(dataset_files) >= num_hosts:
              # Sharding on data sources (e.g. filenames) if there are enough files.
              dataset_files = np.array_split(dataset_files, num_hosts)[
                  host_id
              ]
              dataset = tf.data.TFRecordDataset(
                  dataset_files,
                  num_parallel_reads=tf.data.experimental.AUTOTUNE,
              )
          else:
              # Sharding using tf.shard.
              dataset = tf.data.TFRecordDataset(
                  dataset_files,
                  num_parallel_reads=tf.data.experimental.AUTOTUNE,
              )
              dataset = dataset.shard(
                  num_shards=num_hosts, index=host_id
              )
  
          # pylint: disable=g-long-lambda
          dataset = dataset.map(
              lambda x: pileup_parser(
                  x, mode="train", pileup_height=299, pileup_width=299
              ),
              num_parallel_calls=tf.data.experimental.AUTOTUNE,
          )
          dataset = dataset.map(
              pileup_normalizer,
              num_parallel_calls=tf.data.experimental.AUTOTUNE,
          )
  
          dataset = dataset.map(
              preprocess,
              num_parallel_calls=tf.data.experimental.AUTOTUNE,
          )
          return dataset
  
      # Use different seed for each host.
      if shuffle_seed is None:
          local_seed = None
      else:
          data_seed = 0
          local_seed = data_seed + jax.process_index()
  
      train_dataset = build_dataset(
          dataset_fn=build_pileup_window_dataset,
          batch_size=batch_size,
          seed=local_seed,
          split="train",
          strategy=None,
      )
  
      if dataset_service_address:
          if shuffle_seed is not None:
              raise ValueError(
                  "Using dataset service with a random seed causes each "
                  "worker to produce exactly the same data. Add "
                  "config.shuffle_seed = None to your config if you "
                  "want to run with dataset service."
              )
          logging.info(
              "Using the tf.data service at %s", dataset_service_address
          )
          train_dataset = dataset_utils.distribute(
              train_dataset, dataset_service_address
          )
  
      eval_dataset = build_dataset(
          dataset_fn=build_pileup_window_dataset,
          split="valid",
          batch_size=eval_batch_size,
          strategy=None,
      )
  
      test_dataset = build_dataset(
          dataset_fn=build_pileup_window_dataset,
          split="test",
          batch_size=eval_batch_size,
          strategy=None,
      )
  
      shard_batches = functools.partial(
          dataset_utils.shard, n_devices=num_shards
      )
      maybe_pad_batches_train = functools.partial(
          dataset_utils.maybe_pad_batch,
          train=True,
          batch_size=batch_size,
      )
      maybe_pad_batches_eval = functools.partial(
          dataset_utils.maybe_pad_batch,
          train=False,
          batch_size=eval_batch_size,
      )
  
      train_iter = iter(train_dataset)
      train_iter = map(dataset_utils.tf_to_numpy, train_iter)
      train_iter = map(maybe_pad_batches_train, train_iter)
      train_iter = map(shard_batches, train_iter)
      train_iter = jax_utils.prefetch_to_device(
          train_iter, prefetch_buffer_size
      )
  
      valid_iter = iter(eval_dataset)
      valid_iter = map(dataset_utils.tf_to_numpy, valid_iter)
      valid_iter = map(maybe_pad_batches_eval, valid_iter)
      valid_iter = map(shard_batches, valid_iter)
      valid_iter = jax_utils.prefetch_to_device(
          valid_iter, prefetch_buffer_size
      )
  
      test_iter = iter(test_dataset)
      test_iter = map(dataset_utils.tf_to_numpy, test_iter)
      test_iter = map(maybe_pad_batches_eval, test_iter)
      test_iter = map(shard_batches, test_iter)
      test_iter = jax_utils.prefetch_to_device(
          test_iter, prefetch_buffer_size
      )
  
      num_classes = 3
      image_size = 256
      input_shape = [-1, image_size, image_size, 7]
  
      meta_data = {
          "num_classes": num_classes,
          "input_shape": input_shape,
          "num_train_examples": 30_000 * 19,
          "num_eval_examples": 30_000 * 5,
          "num_test_examples": 30_000,
          "test_name": get_dataset_name(dataset_configs.test_path),
          "eval_name": get_dataset_name(dataset_configs.eval_path),
          "input_dtype": getattr(jnp, dtype_str),
          "target_is_onehot": True,
      }
  
      return dataset_utils.Dataset(
          train_iter, valid_iter, test_iter, meta_data
      )
  
file_path: scenic/scenic/projects/svvit/datasets/pileup_window_dataset.py
imports: import functools
  from typing import Optional, Dict, Tuple
  from absl import logging
  from flax import jax_utils
  import jax
  import jax.numpy
  import numpy
  from scenic.dataset_lib import dataset_utils
  from scenic.dataset_lib import datasets
  import tensorflow
label_desc: Use image node_id 7 for calling a global function with example usage: get_dataset() and returns: dataset_utils, dataset