def evaluate(model: nn.Module,
             state: TrainState,
             eval_ds: tf.data.Dataset,
             num_eval_steps: int = -1) -> Union[None, EvalMetrics]:
  """Evaluate the model on the given dataset."""
  logging.info("Starting evaluation.")
  eval_metrics = None
  with StepTraceContextHelper("eval", 0) as trace_context:
    for step, batch in enumerate(eval_ds):  # pytype: disable=wrong-arg-types
      batch = jax.tree_map(np.asarray, batch)
      metrics_update = flax_utils.unreplicate(eval_step(model, state, batch))
      eval_metrics = (
          metrics_update
          if eval_metrics is None else eval_metrics.merge(metrics_update))
      if num_eval_steps > 0 and step + 1 == num_eval_steps:
        break
      trace_context.next_step()
  return eval_metrics