def generate_masked_tokens(batch: Batch,
                           rng: Optional[jnp.ndarray],
                           mask_ratio: jnp.ndarray,
                           task: task_registry.Task,
                           config: ml_collections.ConfigDict,
                           *,
                           no_latent: bool = False):
  """Generate masked input tokens with condition tokens.

  Args:
    batch: A single batch of data, where batch['inputs'] is latent code of shape
      [device_bs, l_t, l_h, l_w] in jnp.int32, optionally batch['cond_inputs']
      is latent code like batch['inputs'] and batch['label'] is class id of
      shape [device_bs, 1] in jnp.int32.
    rng: Optional PRNGKey for task.mask_fn.
    mask_ratio: of shape [1] or shape [device_bs] where values are in (0,1].
    task: current task.
    config: model config.
    no_latent: whether only returning prefix tokens without latent tokens.

  Returns:
    Dictionary of the masked and target token sequence with the class info.
  """
  batch_size = batch['inputs'].shape[0]
  tokens_list = []

  # Task condition token
  if config.get('multi_task', False):
    task_tokens = jnp.full((batch_size, 1), task.task_token, jnp.int32)
    if any([task.startswith('full_generation') for task in config.tasks]):
      uncond_task_tokens = jnp.full(
          (batch_size, 1), task_registry.SpecialTokens.TASK_FULL_GENERATION,
          jnp.int32)
    else:
      uncond_task_tokens = task_tokens
    tokens_list.append(
        mask_utils.no_mask({
            'inputs': task_tokens,
            'uncond_inputs': uncond_task_tokens
        }))

  # Class condition token
  class_tokens = get_class_tokens(batch, config, task.config)
  uncond_class_tokens = get_class_tokens(
      batch, config, task.config, return_uncond_token=True)
  tokens_list.append(
      mask_utils.no_mask({
          'inputs': class_tokens,
          'uncond_inputs': uncond_class_tokens
      }))

  # Frame condition tokens, deprecated, only for single frame prediction task
  if (task.task_id == 'frame_prediction' and
      task.config.get('condition_mode') == 'prefix'):
    num_frames = config.frame_prediction.cond_latent_frames
    cond_latents = batch.pop('cond_inputs')[:, :num_frames]
    tokens_list.append(mask_utils.no_mask({'inputs': cond_latents}))
    sep_token_ids = jnp.full((batch_size, 1), task_registry.SpecialTokens.SEP,
                             jnp.int32)
    tokens_list.append(mask_utils.no_mask({'inputs': sep_token_ids}))
    batch.pop('cond_mask')

  # Latent tokens
  if not no_latent:
    latent_tokens = task.mask_fn(batch, rng, mask_ratio)
    tokens_list.append(latent_tokens)

  all_tokens = mask_utils.concat_tokens(*tokens_list)

  return all_tokens
Generate masked input tokens with condition tokens.

  Args:
    batch: A single batch of data, where batch['inputs'] is latent code of shape
      [device_bs, l_t, l_h, l_w] in jnp.int32, optionally batch['cond_inputs']
      is latent code like batch['inputs'] and batch['label'] is class id of
      shape [device_bs, 1] in jnp.int32.
    rng: Optional PRNGKey for task.mask_fn.
    mask_ratio: of shape [1] or shape [device_bs] where values are in (0,1].
    task: current task.
    config: model config.
    no_latent: whether only returning prefix tokens without latent tokens.

  Returns:
    Dictionary of the masked and target token sequence with the class info.
  
videogvt/train_lib/maskgvt_train_lib.py