code: def get_dataset(
      *,
      batch_size,
      eval_batch_size,
      num_shards,
      dtype_str="float32",
      shuffle_seed=0,
      prefetch_buffer_size=2,
      rng=None,
      dataset_configs=None,
      dataset_service_address: Optional[str] = None,
  ):
      """Returns generators for the Cityscapes 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.
        prefetch_buffer_size: int; Buffer size for the TFDS prefetch.
        rng: JAX rng key, which can be used for augmentation, shuffling, etc.
        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 batch_size
      del shuffle_seed, rng
      del dataset_service_address
  
      dtype = getattr(tf, dtype_str)
      dataset_configs = dataset_configs or {}
      dataset_name = _check_dataset_exists(dataset_configs)
      cityscapes_variant_info = DATASET_INFO.get(dataset_name)
      target_size = dataset_configs.get("target_size", None)
  
      if "cityscapes" in dataset_name:
          preprocess_example = preprocess_examples["cityscapes"]
      elif "fishyscapes" in dataset_name:
          preprocess_example = preprocess_examples["fishyscapes"]
  
      preprocess_ex_eval = functools.partial(
          preprocess_example,
          train=False,
          dtype=dtype,
          resize=target_size,
      )
  
      logging.info(
          "Loading validation split of the %s dataset.", dataset_name
      )
  
      eval_ds, _ = cityscapes_load_split(
          dataset_name=dataset_name,
          batch_size=eval_batch_size,
          train=False,
          dtype=dtype,
          preprocess_ex_eval=preprocess_ex_eval,
      )
  
      maybe_pad_batches_eval = functools.partial(
          dataset_utils.maybe_pad_batch,
          train=False,
          batch_size=eval_batch_size,
          pixel_level=True,
      )
      shard_batches = functools.partial(
          dataset_utils.shard, n_devices=num_shards
      )
      exclude_classes = functools.partial(
          cityscapes_dataset.exclude_bad_classes,
          new_labels=cityscapes_dataset.get_post_exclusion_labels(),
      )
  
      eval_iter = iter(eval_ds)
      eval_iter = map(dataset_utils.tf_to_numpy, eval_iter)
      eval_iter = map(maybe_pad_batches_eval, eval_iter)
  
      if "cityscapes" in dataset_name:
          eval_iter = map(exclude_classes, eval_iter)
      eval_iter = map(shard_batches, eval_iter)
      eval_iter = jax_utils.prefetch_to_device(
          eval_iter, prefetch_buffer_size
      )
  
      if target_size is None:
          input_shape = (-1, 1024, 2048, 3)
      else:
          input_shape = (-1,) + tuple(target_size) + (3,)
  
      meta_data = {
          "input_shape": input_shape,
          "num_train_examples": 0,
          "num_eval_examples": cityscapes_variant_info[
              "num_of_examples"
          ],
          "input_dtype": getattr(jnp, dtype_str),
          "target_is_onehot": False,
      }
  
      if "cityscapes" in dataset_name:
          meta_data.update(cityscapes_meta_data)
      elif "fishyscapes" in dataset_name:
          meta_data.update(fishyscapes_meta_data)
  
      return dataset_utils.Dataset(None, eval_iter, None, meta_data)
  
file_path: scenic/scenic/projects/robust_segvit/datasets/cityscapes_variants.py
imports: import functools
  from typing import Optional
  from absl import logging
  from flax import jax_utils
  import jax.numpy
  from scenic.dataset_lib import cityscapes_dataset
  from scenic.dataset_lib import dataset_utils
  from scenic.dataset_lib import datasets
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 5 for calling a global function with example usage: get_dataset() and returns: dataset_utils