def eval_step(
    train_state,
    model_vars,
    batch: Dict[str, Any],
    model_config: ml_collections.FrozenConfigDict,
) -> Tuple[Dict[str, Any], Dict[str, Any]]:
  """Calculate evaluation metrics on a batch.

  Args:
    train_state: contains model params, loss fn, grad update fn.
    model_vars: model variables that are not optimized.
    batch: input to model.
    model_config: contains model hyperparameters.

  Returns:
    Dictionary of metrics and auxiliary output.
  """

  def eval_step_partial(model_params):
    return train_state.apply_fn(
        model_config,
        model_params,
        model_vars,
        batch,
        deterministic=True,
    )

  _, metrics, auxiliary_output = eval_step_partial(train_state.params)
  metrics = metric_utils.update_metrics_dtype(metrics)
  metrics = jax.lax.psum(metrics, axis_name='batch')

  return metrics, auxiliary_output
Calculate evaluation metrics on a batch.

  Args:
    train_state: contains model params, loss fn, grad update fn.
    model_vars: model variables that are not optimized.
    batch: input to model.
    model_config: contains model hyperparameters.

  Returns:
    Dictionary of metrics and auxiliary output.
  
language/mentionmemory/training/trainer.py