def _compute_stats(
    params, batch, model,
    scoring_fn):
  """Runs inference and computes model predictions.

  Args:
    params: Model state (parameters).
    batch: Current batch of examples.
    model: The model itself. Flax separates model state and architecture.
    scoring_fn: Task dependent function mapping raw model output logits to
      prediction.

  Returns:
    Model predictions along with example labels.
  """
  inputs = {
      "input_ids": batch["input_ids"],
      "type_ids": batch["type_ids"],
      "deterministic": True
  }

  y = model.apply({"params": params}, **inputs)
  result = {
      "idx": batch["idx"],
      "label": batch["label"],
      "prediction": scoring_fn(y),
      "input_ids": batch["input_ids"],  # Required for SQuAD F1 metric
  }

  return result