def train_step(
    train_state: TrainState, batch: Batch, task: Task, *,
    flax_model: nn.Module,
    config: ml_collections.ConfigDict) -> Tuple[Any, metrics.Collection]:
  """Runs a single step of training.

  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
    task: Current training task.
    flax_model: Flax model
    config: Configurations of the experiment.

  Returns:
    Updated state of training, computed metrics.
  """
  new_rng, ratio_rng, mask_rng, dropout_rng = jax.random.split(
      train_state.rng, 4)
  vq_codebook_size = maskgvt_train_lib.get_vq_codebook_size(config)

  # TODO(Lijun-Yu): try different mask_ratio inside a batch
  mask_ratio = maskgvt_train_lib.generate_mask_ratio(ratio_rng, config)
  batch_tokens = maskgvt_train_lib.generate_masked_tokens(
      batch, mask_rng, mask_ratio, task, config)

  # Bind the dropout rng to the host/device we are on.
  dropout_rng = scenic_train_utils.bind_rng_to_host_device(
      dropout_rng, axis_name='device', bind_to='device')
  mutable = ['batch_stats', 'spectral_norm_stats']

  def training_loss_fn(params):
    variables = {'params': params, **train_state.model_state}
    logits, new_model_state = flax_model.apply(
        variables,
        batch_tokens['masked_inputs'],
        batch_tokens['segment_ids'],
        deterministic=False,
        mutable=mutable,
        rngs={'dropout': dropout_rng})
    # logits shape [bs, 1 + (l_cond) + l_t * l_h * l_w,
    #               vq_codebook_size + num_classes + num_special_tokens]
    logits = logits[:, :, :vq_codebook_size]
    # TODO(roadjiang): add classification loss. Use batch_tokens['weights'][0]
    one_hot_targets = jax.nn.one_hot(batch_tokens['targets'], logits.shape[-1])

    sof_ce_loss = (
        scenic_model_utils.weighted_unnormalized_softmax_cross_entropy(
            logits,
            one_hot_targets,
            weights=batch.get('batch_mask'),
            label_smoothing=config.get('label_smoothing'),
        )
    )

    weights = batch_tokens['weights']
    masked_sof_ce_loss = jnp.sum(
        sof_ce_loss * weights, axis=-1) / (
            jnp.sum(weights, axis=-1) + 1e-8)
    token_loss = jnp.mean(masked_sof_ce_loss)

    l2_loss = 0
    if config.get('l2_decay_factor') is None:
      total_loss = token_loss
    else:
      l2_loss = scenic_model_utils.l2_regularization(params)
      total_loss = token_loss + 0.5 * config.l2_decay_factor * l2_loss

    return total_loss, (new_model_state, token_loss, l2_loss)

  compute_gradient_fn = jax.value_and_grad(training_loss_fn, has_aux=True)
  (total_loss, (new_model_state, token_loss,
                l2_loss)), grad = compute_gradient_fn(train_state.params)

  # Re-use same axis_name as in the call to `pmap(...train_step...)` below.
  grad = jax.lax.pmean(grad, axis_name='device')
  if config.get('max_grad_norm') is not None:
    grad = clip_grads(grad, config.max_grad_norm)
  grad_norm = l2_norm(grad)

  updates, new_opt_state = train_state.tx.update(grad, train_state.opt_state,
                                                 train_state.params)
  new_params = optax.apply_updates(train_state.params, updates)

  new_train_state = train_state.replace(  # pytype: disable=attribute-error
      global_step=train_state.global_step + 1,
      opt_state=new_opt_state,
      params=new_params,
      model_state=new_model_state,
      rng=new_rng)
  metrics_update = TrainMetrics.gather_from_model_output(
      axis_name='device',
      total_loss=total_loss,
      token_loss=token_loss,
      l2_loss=l2_loss,
      mask_ratio=jnp.mean(mask_ratio),
      grad_norm=grad_norm)
  return new_train_state, metrics_update
Runs a single step of training.

  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
    task: Current training task.
    flax_model: Flax model
    config: Configurations of the experiment.

  Returns:
    Updated state of training, computed metrics.
  
videogvt/trainers/maskgvt_trainer.py