def _get_scenic_dataset(
    config: ml_collections.ConfigDict,
    data_rng: PRNGKey,
    batch_size: int,
    *,
    start_step: Optional[int] = None,
    num_local_shards: Optional[int] = None,
    dataset_service_address: Optional[str] = None,
    dataset_name: Optional[str] = None,
    dataset_configs: Optional[ml_collections.ConfigDict] = None
) -> scenic_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.
    batch_size: Global batch size which applies to all data splits.
    start_step: Start step for GRAIN-backed inpute pipelines.
    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.

  Returns:
    A dataset_utils.Dataset object.
  """
  dataset_name = dataset_name or config.dataset_name
  if dataset_name == 'unit_test':
    logging.warn('----->>Using the pseudo dataset. Used only in unit test!')
    return get_pseudo_dataset(config, data_rng, batch_size)

  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())

  if dataset_name in _CUSTOM_DATASET_TABLE:
    module = _CUSTOM_DATASET_TABLE[dataset_name]
    logging.info('On-demand import of dataset (%s) from module (%s).',
                 dataset_name, module)
    dataset_builder = importlib.import_module(module).get_dataset
  else:
    dataset_builder = scenic_datasets.get_dataset(dataset_name)

  if batch_size % device_count > 0:
    raise ValueError(f'Batch size ({batch_size}) must be divisible by the '
                     f'number of devices ({device_count})')

  local_batch_size = 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',
                                                  ml_collections.ConfigDict())
  num_local_shards = num_local_shards or jax.local_device_count()


  if dataset_configs.get('expects_start_step'):
    kwargs = {'start_step': start_step}
  else:
    kwargs = {}

  dataset = dataset_builder(
      batch_size=local_batch_size,
      eval_batch_size=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
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.
    batch_size: Global batch size which applies to all data splits.
    start_step: Start step for GRAIN-backed inpute pipelines.
    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.

  Returns:
    A dataset_utils.Dataset object.
  
videogvt/train_lib/train_utils.py