def all_mask(batch: TokenDict,
             rng: Optional[jnp.ndarray] = None,
             mask_ratio: Optional[jnp.ndarray] = None,
             *,
             weight: float = 1.) -> TokenDict:
  """Mask all tokens, except for the condition tokens."""
  del rng, mask_ratio
  if batch.get('cond_inputs') is not None:
    masked_inputs = jnp.where(batch['cond_mask'], batch['cond_inputs'],
                              MASK_TOKEN)
  else:
    masked_inputs = jnp.full_like(batch['inputs'], MASK_TOKEN)
  uncond_masked_inputs = jnp.full_like(batch['inputs'], MASK_TOKEN)
  return {
      'masked_inputs': masked_inputs,
      'uncond_masked_inputs': uncond_masked_inputs,
      'segment_ids': jnp.zeros_like(batch['inputs']),
      'targets': batch['inputs'],
      'weights': jnp.full_like(masked_inputs, weight, jnp.float32),
  }
Mask all tokens, except for the condition tokens.
videogvt/train_lib/mask_utils.py