def evaluate(
    eval_step_fn,
    train_state: ts.TrainState,
    model_vars: Dict[str, Any],
    eval_data: Sequence[Dict[str, Any]],
) -> Tuple[Dict[str, Any], Sequence[Tuple[Dict[str, Any], Optional[Dict[
    str, Any]]]]]:
  """Evaluate current parameters and return a dictionary with metrics.

  Args:
    eval_step_fn: partial eval step that takes in model params and inputs only
    train_state: contains model params, loss fn, grad update fn.
    model_vars: model variables that are not optimized.
    eval_data: sequence of evaluation data.

  Returns:
    Dictionary of metrics aggregated over all evaluation steps and the info
    for the very first batch (batch itself and corresponding auxiliary output).
  """

  logging.info('Performing evaluation.')
  eval_metrics = []
  eval_auxiliary = []
  for batch in eval_data:
    batch = jax.tree.map(jnp.asarray, batch)
    metrics, auxiliary_output = eval_step_fn(
        train_state,
        model_vars,
        batch,
    )
    eval_metrics.append(metrics)
    batch_auxiliary = (jax.device_get(batch), jax.device_get(auxiliary_output))
    eval_auxiliary.append(batch_auxiliary)
  eval_metrics = common_utils.get_metrics(eval_metrics)
  eval_metrics_sums = jax.tree.map(jnp.sum, eval_metrics)
  eval_summary = metric_utils.process_metrics(eval_metrics_sums, prefix='eval')
  return eval_summary, eval_auxiliary
Evaluate current parameters and return a dictionary with metrics.

  Args:
    eval_step_fn: partial eval step that takes in model params and inputs only
    train_state: contains model params, loss fn, grad update fn.
    model_vars: model variables that are not optimized.
    eval_data: sequence of evaluation data.

  Returns:
    Dictionary of metrics aggregated over all evaluation steps and the info
    for the very first batch (batch itself and corresponding auxiliary output).
  
language/mentionmemory/training/trainer.py