def sample_step(train_state: TrainState, batch: Batch, task: Task,
                *, flax_model: nn.Module,
                config: ml_collections.ConfigDict):
  """Runs a single step to predict tokens during training.

  Args:
    train_state: The state of training TrainState
    batch: A single batch of data. Dictionary where
      batch['input_indices'].shape= device_bs, l_t*l_h*l_w+1, vocabulary_size
    task: Current training task.
    flax_model: Flax model
    config: Configurations of the experiment.

  Returns:
    Token predition visualized as images formatted in spatial grids.
  """
  vq_codebook_size = maskgvt_train_lib.get_vq_codebook_size(config)
  batch_tokens = maskgvt_train_lib.generate_tokens_for_lm(batch, task, config)
  variables = {'params': train_state.params, **train_state.model_state}
  logits = flax_model.apply(
      variables, batch_tokens['inputs'], deterministic=True)
  # only illustrate the last latent_seq_len tokens.
  logits = logits[..., :vq_codebook_size]
  # we do not do sampling.sampling as there is no sampling during training
  # do not use batch_tokens['inputs'] here as it has mask token as bos.
  pred_visualization = maskgvt_train_lib.illustrate_predictions(
      batch_tokens['targets'], logits, batch_tokens['targets'], config)
  outputs = maskgvt_train_lib.draw_video_boundary(pred_visualization)
  results = {}
  return outputs, results
Runs a single step to predict tokens during training.

  Args:
    train_state: The state of training TrainState
    batch: A single batch of data. Dictionary where
      batch['input_indices'].shape= device_bs, l_t*l_h*l_w+1, vocabulary_size
    task: Current training task.
    flax_model: Flax model
    config: Configurations of the experiment.

  Returns:
    Token predition visualized as images formatted in spatial grids.
  
videogvt/trainers/lmgvt_trainer.py