code: def get_datasets(
      config: ml_collections.ConfigDict,
      data_rng: jnp.ndarray,
      dataset_service_address: Optional[str] = None,
  ):
      """Creates dataset from config."""
  
      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_dict = {}
      for ds_name, cfg in config.datasets.items():
          # This key is needed for disabling datasets in hyperparameter sweeps.
          if cfg.get("dont_use") is not None and cfg.get("dont_use"):
              continue
  
          if config.get("batch_sizes") is not None:
              batch_size = config.batch_sizes.get(ds_name)
          else:
              batch_size = config.batch_size
  
          if batch_size % device_count > 0:
              raise ValueError(
                  f"Batch size ({batch_size}) of {ds_name} must be divisible "
                  f"by the number of devices ({device_count})"
              )
  
          if config.get("eval_batch_sizes") is not None:
              eval_batch_size = config.eval_batch_sizes.get(ds_name)
          else:
              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}) of {ds_name} must be "
                  f"divisible 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 of %s : %d", ds_name, local_batch_size
          )
          logging.info(
              "device_batch_size of %s : %d", ds_name, device_batch_size
          )
  
          shuffle_seed = cfg.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."
              )
  
          # 'bit' consists of many datasets, so we do this to have a unique dataset
          # key if we train on multiple datasets from 'bit'. E.g. ds_name =
          # 'bit_caltech'.
          if ds_name.startswith("bit_"):
              dataset_builder = datasets.get_dataset("bit")
          elif ds_name in ["kinetics400", "moments_in_time"]:
              dataset_builder = datasets.get_dataset(
                  "video_tfrecord_dataset"
              )
          elif ds_name in ["vggsound", "audioset"]:
              dataset_builder = datasets.get_dataset(
                  "audiovisual_tfrecord_dataset"
              )
          else:
              dataset_builder = datasets.get_dataset(ds_name)
  
          dataset_rng, data_rng = jax.random.split(data_rng)
          ds = dataset_builder(
              batch_size=local_batch_size,
              eval_batch_size=eval_local_batch_size,
              num_shards=jax.local_device_count(),
              dtype_str=cfg.data_dtype_str,
              rng=dataset_rng,
              shuffle_seed=shuffle_seed,
              dataset_configs=cfg,
              dataset_service_address=dataset_service_address,
          )
  
          # Add task information to the dataset meta_data:
          ds.meta_data["task"] = cfg.task
          ds.meta_data["modality"] = cfg.get("modality", "image")
          dataset_dict[ds_name] = ds
  
      return dataset_dict
  
file_path: scenic/scenic/projects/polyvit/train_utils.py
imports: from typing import Any, Dict, Optional, Tuple
  from absl import logging
  from flax.training import checkpoints
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.dataset_lib import datasets
  from scenic.train_lib import train_utils
  from tensorflow.io import gfile
label_desc: Use image node_id 3 for calling a global function with example usage: get_datasets(config, data_rng, dataset_service_address) and returns: dataset_dict