def get_class_tokens(batch: Batch, config: ml_collections.ConfigDict,
                     task_config: ml_collections.ConfigDict,
                     return_uncond_token: bool = False):
  """Gets the class tokens for a batch."""
  offset = get_vq_codebook_size(config)
  if task_config.get('class_conditional', False) and not return_uncond_token:
    class_token = batch['label'] + offset
  else:
    uncond_token = config.dataset_configs.num_classes + offset
    class_token = jnp.full((batch['inputs'].shape[0], 1), uncond_token,
                           jnp.int32)
  return class_token
Gets the class tokens for a batch.
videogvt/train_lib/maskgvt_train_lib.py