def generate_step(train_state: TrainState,
                  batch: Batch,
                  cache: Dict[str, jnp.ndarray],
                  task: Task,
                  rng: jnp.ndarray,
                  *,
                  flax_model: nn.Module,
                  tokenizer_dict: Dict[str, TokenizerFn],
                  config: ml_collections.ConfigDict):
  """Runs generation given a batch of inputs.

  Args:
    train_state: The state of training TrainState
    batch: A single batch of data. Dictionary where batch['inputs'].shape=
      device_bs, t, h, w, 3 batch['batch_mask'].shape= device_bs
      batch['batch_mask'].dtype= float32 where batch['batch_mask'] > 0 for valid
      examples
    cache: Intermediate variables for self-attention layers.
    task: Current generation task.
    rng: The PRNGKey for generation sampling
    flax_model: Flax model
    tokenizer_dict: A dictionary of tokenizer and detokenizer
    config: Configurations of the experiment.

  Returns:
    Sampled videos formatted in spatial grids.
  """
  variables = {'params': train_state.params, **train_state.model_state}
  latent_shape = (batch['inputs'].shape[0], *config.transformer.latent_shape)

  # Tokenize
  input_video = batch.pop('inputs')
  outputs = {
      'original_video': input_video,
      'batch_mask': batch['batch_mask']
  }
  if task.cond_fn is not None:
    cond_dict = task.cond_fn(input_video)
    outputs['condition_video'] = cond_dict['video']
    outputs['condition_mask'] = cond_dict['video_mask']
    batch['cond_inputs'] = tokenizer_dict['tokenizer'](cond_dict.pop('video'))
    batch['cond_mask'] = cond_dict.pop('latent_mask')
    outputs['condition_latents'] = batch['cond_inputs']
  batch['inputs'] = jnp.zeros(latent_shape, dtype=jnp.int32)

  if 'label' in batch:
    outputs['label'] = batch['label']

  lm_tokens = maskgvt_train_lib.generate_tokens_for_lm(
      batch, task, config)
  inputs, cond_lengths = lm_tokens['inputs'], lm_tokens['cond_lengths']
  generated_seq = maskgvt_train_lib.lm_decode(
      inputs, cond_lengths, cache, rng, variables, flax_model, config)
  outputs['generated_latents'] = generated_seq.reshape(latent_shape)
  # Detokenize
  generated_video = tokenizer_dict['detokenizer'](outputs['generated_latents'])
  if task.cond_fn is not None and config.sampling.get(
      'replace_condition_pixels', False):
    generated_video = jnp.where(
        cond_dict.pop('video_mask'), input_video, generated_video)
  outputs['generated_video'] = jnp.clip(generated_video, 0, 1)
  results = {}
  return outputs, results
Runs generation given a batch of inputs.

  Args:
    train_state: The state of training TrainState
    batch: A single batch of data. Dictionary where batch['inputs'].shape=
      device_bs, t, h, w, 3 batch['batch_mask'].shape= device_bs
      batch['batch_mask'].dtype= float32 where batch['batch_mask'] > 0 for valid
      examples
    cache: Intermediate variables for self-attention layers.
    task: Current generation task.
    rng: The PRNGKey for generation sampling
    flax_model: Flax model
    tokenizer_dict: A dictionary of tokenizer and detokenizer
    config: Configurations of the experiment.

  Returns:
    Sampled videos formatted in spatial grids.
  
videogvt/trainers/lmgvt_trainer.py