def tokenize_with_cond(inputs, cond_fn, *, tokenizer_fn):
  """Run tokenizer on inputs and derived condition."""
  if cond_fn is not None:
    # Use cond latents from masked inputs.
    cond_dict = cond_fn(inputs)
    all_inputs = jnp.concatenate([inputs, cond_dict.pop('video')])
    all_tokens = tokenizer_fn(all_inputs)
    input_tokens, cond_tokens = jnp.split(all_tokens, 2)
    return {
        'inputs': input_tokens,
        'cond_inputs': cond_tokens,
        'cond_mask': cond_dict['latent_mask']
    }
  else:
    input_tokens = tokenizer_fn(inputs)
    return {'inputs': input_tokens}
Run tokenizer on inputs and derived condition.
videogvt/train_lib/maskgvt_train_lib.py