code: def get_dataset(
      *,
      batch_size: int,
      eval_batch_size: int,
      num_shards: int,
      dtype_str: Text = "float32",
      shuffle_seed: Optional[int] = 0,
      rng: Optional[Rng] = None,
      dataset_configs: ml_collections.ConfigDict,
      dataset_service_address: Optional[str] = None,
  ) -> dataset_utils.Dataset:
      """Returns a generator for the dataset."""
      del rng  # Parameter was required by caller API, but is unused.
  
      def validate_config(field):
          if dataset_configs.get(field) is None:
              raise ValueError(
                  f"{field} must be specified for TFRecord dataset."
              )
  
      validate_config("base_dir")
      validate_config("tables")
      validate_config("examples_per_subset")
      validate_config("num_classes")
  
      num_frames = dataset_configs.get("num_frames", 32)
      num_test_clips = dataset_configs.get("num_test_clips", 1)
      stride = dataset_configs.get("stride", 2)
      min_resize = dataset_configs.get("min_resize", 256)
      crop_size = dataset_configs.get("crop_size", 224)
      one_hot_label = dataset_configs.get("one_hot_label", True)
      zero_centre_data = dataset_configs.get("zero_centering", True)
      augmentation_params = dataset_configs.get(
          "augmentation_params", None
      )
      num_train_val_clips = dataset_configs.get(
          "num_train_val_clips", 1
      )
      keep_test_key = dataset_configs.get("keep_test_key", False)
      # For the test set, the actual batch size is
      # test_batch_size * num_test_clips.
      test_batch_size = dataset_configs.get(
          "test_batch_size", eval_batch_size
      )
      do_three_spatial_crops = dataset_configs.get(
          "do_three_spatial_crops", False
      )
      num_spatial_crops = 3 if do_three_spatial_crops else 1
      test_split = dataset_configs.get("test_split", "test")
      label_offset = dataset_configs.get("label_offset", 0)
  
      ds_factory = functools.partial(
          TFRecordDatasetFactory,
          base_dir=dataset_configs.base_dir,
          tables=dataset_configs.tables,
          examples_per_subset=dataset_configs.examples_per_subset,
          num_classes=dataset_configs.num_classes,
          num_groups=jax.process_count(),
          group_index=jax.process_index(),
      )
  
      def create_dataset_iterator(
          subset: Text,
          batch_size_local: int,
          num_clips: int,
          keep_key_local: bool = False,
      ) -> Tuple[Iterator[Batch], int]:
          is_training = subset == "train"
          is_test = subset == "test"
          logging.info("Loading split %s", subset)
  
          dataset, num_examples = load_split(
              ds_factory,
              batch_size=batch_size_local,
              subset=subset,
              num_frames=num_frames,
              stride=stride,
              num_test_clips=num_clips,
              min_resize=min_resize,
              crop_size=crop_size,
              one_hot_label=one_hot_label,
              zero_centering=zero_centre_data,
              augmentation_params=augmentation_params,
              keep_key=keep_key_local,
              do_three_spatial_crops=do_three_spatial_crops and is_test,
              label_offset=label_offset,
          )
  
          if dataset_service_address and is_training:
              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,
              )
              dataset = dataset_utils.distribute(
                  dataset, dataset_service_address
              )
  
          pad_batch_size = batch_size_local
          if is_test:
              pad_batch_size = (
                  batch_size_local * num_clips * num_spatial_crops
              )
          maybe_pad_batches = functools.partial(
              dataset_utils.maybe_pad_batch,
              train=is_training,
              batch_size=pad_batch_size,
          )
          shard_batches = functools.partial(
              dataset_utils.shard, n_devices=num_shards
          )
  
          current_ds_iterator = (
              map_keys(dataset_utils.tf_to_numpy(data))
              for data in iter(dataset)
          )
  
          if is_test and num_clips * num_spatial_crops > 1:
              current_ds_iterator = map(
                  tile_label_key, current_ds_iterator
              )
  
          current_ds_iterator = map(
              maybe_pad_batches, current_ds_iterator
          )
          current_ds_iterator = map(shard_batches, current_ds_iterator)
          if is_training and dataset_configs.get("prefetch_to_device"):
              # Async bind batch to device which speeds up training.
              current_ds_iterator = jax_utils.prefetch_to_device(
                  current_ds_iterator,
                  dataset_configs.get("prefetch_to_device"),
              )
  
          return current_ds_iterator, num_examples
  
      train_iter, n_train_examples = create_dataset_iterator(
          "train", batch_size, num_train_val_clips
      )
      eval_iter, n_eval_examples = create_dataset_iterator(
          "validation", eval_batch_size, num_train_val_clips
      )
      test_iter, n_test_examples = create_dataset_iterator(
          test_split, test_batch_size, num_test_clips, keep_test_key
      )
  
      meta_data = {
          "num_classes": dataset_configs.num_classes,
          "input_shape": (-1, num_frames, crop_size, crop_size, 3),
          "num_train_examples": (
              n_train_examples * num_train_val_clips
          ),
          "num_eval_examples": (n_eval_examples * num_train_val_clips),
          "num_test_examples": (
              n_test_examples * num_test_clips * num_spatial_crops
          ),
          "input_dtype": getattr(jnp, dtype_str),
          "target_is_onehot": True,
      }
      logging.info("Dataset metadata:\n%s", meta_data)
  
      return dataset_utils.Dataset(
          train_iter, eval_iter, test_iter, meta_data
      )
  
imports: import functools
  import os
  from typing import Dict, Iterator, List, Optional, Text, Tuple, Union
  from absl import logging
  from dmvr import builders
  from dmvr import modalities
  from dmvr import video_dataset
  from flax import jax_utils
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.dataset_lib import dataset_utils
  from scenic.dataset_lib import datasets
  from scenic.dataset_lib import video_ops
  from scenic.projects.vivit.data import file_utils
  import tensorflow
label_desc: Use image node_id 5 for calling a global function with example usage: get_dataset() and returns: dataset_utils, current_ds_iterator,  num_examples