def generate_step(train_state: TrainState, batch: Batch, 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
    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']
  # All tokens will be masked so no need to run tokenizer.
  batch['inputs'] = jnp.zeros(latent_shape, dtype=jnp.int32)
  if 'label' in batch:
    outputs['label'] = batch['label']
  # Mask
  rng, subrng = jax.random.split(rng)
  prefix_token_dict = maskgvt_train_lib.generate_masked_tokens(
      batch, subrng, jnp.ones((1,)), task, config,
      no_latent=True)
  # Decode
  rng, subrng = jax.random.split(rng)

  generated_seq = maskgvt_train_lib.fast_decode(batch, prefix_token_dict,  # pytype: disable=wrong-arg-types  # jax-ndarray
                                                task.mask_fn, subrng, variables,
                                                flax_model, None, 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
    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/maskgvt_trainer.py