def eval_step(
    params, batch,
    metric_fn
):
  """Performs a single evaluation step.

  We use this wrapper for parallelizing the evaluation computations across
  multiple devices.

  Args:
    params: Current model state.
    batch: Current batch of evaluation examples.
    metric_fn: Function that maps the model state and batch to output model
      metrics.

  Returns:
    Model metrics for given inputs.
  """
  return metric_fn(params, batch)