code: def cityscapes_load_split(
      dataset_name,
      batch_size,
      train=False,
      dtype=tf.float32,
      shuffle_buffer_size=10,
      shuffle_seed=None,
      data_augmentations=None,
      preprocess_ex_eval=None,
      cache=True,
      data_dir: Optional[str] = None,
  ):
      """Creates a split from the Cityscapes dataset using TensorFlow Datasets.
  
      For the training set, we drop the last partial batch. This is fine to do
      because we additionally shuffle the data randomly each epoch, thus the trainer
      will see all data in expectation. For the validation set, we pad the final
      batch to the desired batch size.
  
      Args:
        dataset_name: string; Dataset name defined in DATASET_INFO.
        batch_size: int; The batch size returned by the data pipeline.
        train: bool; Whether to load the train or evaluation split.
        dtype: TF data type; Data type of the image.
        shuffle_buffer_size: int; Buffer size for the TFDS prefetch.
        shuffle_seed: The seed to use when shuffling the train split.
        data_augmentations: list(str); Types of data augmentation applied on
        preprocess_ex_eval: preprocessing function. Default None.
        cache: bool; Whether to cache dataset in memory.
        data_dir: directory with data.
  
      Returns:
        A `tf.data.Dataset`.
      """
      assert not train, "Only evaluation is supported."
      assert dataset_name in DATASET_INFO
      del data_augmentations
      cityscapes_variant_info = DATASET_INFO.get(dataset_name, {})
      split = cityscapes_variant_info[
          "split"
      ]  # only supports validation
  
      # Load the preprocessing function
      if "cityscapes" in cityscapes_variant_info.get("tfds_name"):
          if dataset_name == "cityscapes":
              builder = tfds.builder(dataset_name, dtype=dtype)
          elif "cityscapes_corrupted" in dataset_name:
              if data_dir is None:
                  # pylint: disable=line-too-long
                  data_dir = "gs://ub-ekb/tensorflow_datasets/cityscapes_corrupted/tfrecords/v.0.0"  # pylint: disable=line-too-long
                  # pylint: enable=line-too-long
              builder = tfds.builder(dataset_name, data_dir=data_dir)
      elif "fishyscapes" in cityscapes_variant_info.get("tfds_name"):
          if data_dir is None:
              data_dir = "gs://ub-ekb/tensorflow_datasets/fishyscapes/tfrecords/v.0.0"
          builder = tfds.builder(dataset_name, data_dir=data_dir)
      else:
          raise NotImplementedError(f"{dataset_name} not available")
  
      ds, ds_info = dataset_utils.load_split_from_tfds_builder(
          builder=builder,
          batch_size=batch_size,
          split=split,
          preprocess_example=preprocess_ex_eval,
          shuffle_buffer_size=shuffle_buffer_size,
          shuffle_seed=shuffle_seed,
          cache=cache,
      )
      return ds, ds_info
  
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 3 for calling a global function with example usage: cityscapes_load_split(dataset_name, batch_size, train, dtype, shuffle_buffer_size, shuffle_seed, data_augmentations, preprocess_ex_eval, cache, data_dir) and returns: ds,  ds_info