def _compute_eval_stats(params, batch,
                        model):
  """Computes pre-training task predictions and stats.

  Args:
    params: Model state (parameters).
    batch: Current batch of examples.
    model: The model itself. Flax separates model state and architecture.

  Returns:
    Model predictions and metrics.
  """
  inputs = {
      "input_ids": batch["input_ids"],
      "type_ids": batch["type_ids"],
      "masked_lm_positions": batch["masked_lm_positions"],
      "masked_lm_labels": batch["masked_lm_ids"],
      "masked_lm_weights": batch["masked_lm_weights"],
      "next_sentence_labels": batch["next_sentence_labels"],
      "deterministic": True
  }

  return model.apply({"params": params}, **inputs)