code: def get_dataset(
      config: ml_collections.ConfigDict,
      data_rng: PRNGKey,
      *,
      num_local_shards: Optional[int] = None,
      dataset_service_address: Optional[str] = None,
      dataset_name: Optional[str] = None,
      dataset_configs: Optional[ml_collections.ConfigDict] = None,
      **kwargs: Any,
  ) -> dataset_utils.Dataset:
      """Creates dataset.
  
      By default, the values in the config file are used.
      However, if the optional `dataset_name` and `dataset_configs` are passed,
        those are used instead.
  
      Args:
        config: The configuration of the experiment.
        data_rng: Random number generator key to use for the dataset.
        num_local_shards: Number of shards for each batch. So (bs, ...) becomes
          (num_local_shards, bs//num_local_shards, ...). If not specified, it will
          be number of local devices.
        dataset_service_address: Used when using the tf.data.experimental.service
        dataset_name: Name of dataset to load, if not reading from the config.
        dataset_configs: Configuration of the dataset, if not reading directly from
          the config.
        **kwargs: Keyword arguments passed to the dataset builders.
  
      Returns:
        A dataset_utils.Dataset object.
      """
      device_count = jax.device_count()
      logging.info("device_count: %d", device_count)
      logging.info("num_hosts : %d", jax.process_count())
      logging.info("host_id : %d", jax.process_index())
  
      dataset_name = dataset_name or config.dataset_name
      dataset_builder = datasets.get_dataset(dataset_name)
  
      batch_size = config.batch_size
      if batch_size % device_count > 0:
          raise ValueError(
              f"Batch size ({batch_size}) must be divisible by the "
              f"number of devices ({device_count})"
          )
  
      eval_batch_size = config.get("eval_batch_size", batch_size)
      if eval_batch_size % device_count > 0:
          raise ValueError(
              f"Eval batch size ({eval_batch_size}) must be divisible "
              f"by the number of devices ({device_count})"
          )
  
      local_batch_size = batch_size // jax.process_count()
      eval_local_batch_size = eval_batch_size // jax.process_count()
      device_batch_size = batch_size // device_count
      logging.info("local_batch_size : %d", local_batch_size)
      logging.info("device_batch_size : %d", device_batch_size)
  
      shuffle_seed = config.get("shuffle_seed", None)
      if dataset_service_address and 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."
          )
  
      dataset_configs = dataset_configs or config.get(
          "dataset_configs", {}
      )
      num_local_shards = num_local_shards or jax.local_device_count()
      dataset = dataset_builder(
          batch_size=local_batch_size,
          eval_batch_size=eval_local_batch_size,
          num_shards=num_local_shards,
          dtype_str=config.data_dtype_str,
          rng=data_rng,
          shuffle_seed=shuffle_seed,
          dataset_configs=dataset_configs,
          dataset_service_address=dataset_service_address,
          **kwargs,
      )
  
      return dataset
  
imports: import collections.abc
  import copy
  import functools
  import os
  import re
  import time
  from typing import Any, Callable, Dict, Mapping, Optional, Sequence, Tuple, Union
  from absl import logging
  from clu import metric_writers
  import flax
  from flax import jax_utils
  from flax import struct
  import flax.linen
  from flax.training import checkpoints
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  import optax
  from scenic.common_lib import debug_utils
  from scenic.dataset_lib import dataset_utils
  from scenic.dataset_lib import datasets
  from scenic.train_lib import optimizers
  from tensorflow.io import gfile
label_desc: Use image node_id 3 for calling a global function with example usage: get_dataset(config, data_rng) and returns: dataset