def get_pseudo_dataset(config: ml_collections.ConfigDict,
                       rng: jnp.ndarray,
                       batch_size: Optional[int] = None):
  """For the unit tests to avoid the cns access issues."""
  if not batch_size:
    batch_size = config.batch_size
  train_ds = []
  num_of_devices = jax.local_device_count()
  local_bs = batch_size // num_of_devices
  image_size = config.image_size
  if isinstance(image_size, int):
    height = width = image_size
  else:
    height, width = image_size
  num_frames = config.dataset_configs.num_frames
  num_classes = config.dataset_configs.num_classes
  num_batches = config.debug_num_batches

  for batch_idx in range(num_batches):
    # (batch, T, H, W, C)
    rng, video_rng, label_rng = jax.random.split(rng, 3)
    videos = jax.random.uniform(
        video_rng,
        shape=(num_of_devices, local_bs, num_frames, height, width,
               3)).astype(jnp.float32)
    label = jax.random.choice(label_rng, num_classes,
                              [num_of_devices, local_bs, 1])
    # Media ids should be an array of lists.
    if config.dataset_configs.get('include_video_id', False):
      media_ids = []
      for d_idx in range(num_of_devices):
        lb_ids = [
            f'id_{batch_idx*batch_size + d_idx*local_bs + lb_idx}'.encode()
            for lb_idx in range(local_bs)
        ]
        media_ids.append(lb_ids)
      media_ids = np.array(media_ids)
    else:
      media_ids = None

    batch_mask = jnp.ones_like(label[..., 0])
    # Simulate incomplete last batch that got padded.
    if batch_idx == num_batches - 1 and config.get(
        'debug_last_batch_padded', False
    ):
      assert (
          local_bs >= 2
      ), 'local_bs should be >= 2 for debug_last_batch_padded option'
      batch_mask = batch_mask.at[:, local_bs - 1:].set(0)
    sample = dict(inputs=videos, label=label, batch_mask=batch_mask)
    if media_ids is not None:
      sample['vid'] = media_ids
    train_ds.append(sample)
  meta_data = dict(
      input_dtype=jnp.float32,
      input_shape=[-1, num_frames, height, width, 3],
      num_classes=num_classes,
      num_train_examples=num_batches * config.batch_size,
      num_eval_examples=num_batches * config.batch_size,
      num_test_examples=num_batches * config.batch_size,
      target_is_onehot=False,
      is_pseudo=True)

  dataset = scenic_dataset_utils.Dataset(
      train_iter=itertools.cycle(train_ds),
      valid_iter=itertools.cycle(train_ds),
      test_iter=itertools.cycle(train_ds),
      meta_data=meta_data)
  return dataset
For the unit tests to avoid the cns access issues.
videogvt/train_lib/train_utils.py