def eval_step(model: Any, state: TrainState,
              batch: Dict[str, jnp.ndarray]) -> metrics.Collection:
  """Compute the metrics for the given model in inference mode.

  The model is applied to the inputs with train=False using all devices on the
  host. Afterwards metrics are averaged across *all* devices (of all hosts).

  Args:
    model: Flax module for the model. The apply method must take input images
      and a boolean argument indicating whether to use training or inference
      mode.
    state: Replicate model state.
    batch: Inputs that should be evaluated.

  Returns:
    Dictionary of the replicated metrics.
  """
  logging.info("eval_step(batch=%s)", batch)
  variables = {
      "params": state.optimizer.target,
  }
  variables.update(state.model_state)
  logits = model(train=False).apply(variables, batch["image"], mutable=False)
  loss = jnp.mean(
      losses.cross_entropy_loss(logits=logits, labels=batch["label"]))
  return EvalMetrics.gather_from_model_output(
      logits=logits,
      labels=batch["label"],
      loss=loss,
      mask=batch.get("mask"),
  )