def train_step(
    train_state: ts.TrainState,
    model_vars: Dict[str, Any],
    batch: Dict[str, Any],
    dropout_rng: jnp.ndarray,
    model_config: ml_collections.FrozenConfigDict,
) -> Tuple[ts.TrainState, Dict[str, Any]]:
  """Perform a single training step.

  Args:
    train_state: contains model params, loss fn, grad update fn.
    model_vars: model variables that are not optimized.
    batch: input to model.
    dropout_rng: seed for dropout rng in model.
    model_config: contains model hyperparameters.

  Returns:
    Train state with updated parameters and dictionary of metrics.
  """

  dropout_rng = jax.random.fold_in(dropout_rng, train_state.step)

  def loss_fn_partial(model_params):
    loss, metrics, _ = train_state.apply_fn(
        model_config,
        model_params,
        model_vars,
        batch,
        deterministic=False,
        dropout_rng={'dropout': dropout_rng},
    )
    return loss, metrics

  grad_fn = jax.value_and_grad(loss_fn_partial, has_aux=True)
  (_, metrics), grad = grad_fn(train_state.params)
  grad = jax.lax.pmean(grad, 'batch')
  metrics = jax.lax.psum(metrics, axis_name='batch')
  metrics = metric_utils.update_metrics_dtype(metrics)
  new_train_state = train_state.apply_gradients(grads=grad)
  return new_train_state, metrics
Perform a single training step.

  Args:
    train_state: contains model params, loss fn, grad update fn.
    model_vars: model variables that are not optimized.
    batch: input to model.
    dropout_rng: seed for dropout rng in model.
    model_config: contains model hyperparameters.

  Returns:
    Train state with updated parameters and dictionary of metrics.
  
language/mentionmemory/training/trainer.py