examples/research_projects/performer/run_mlm_performer.py/compute_metrics
def compute_metrics(logits, labels, weights, label_smoothing=0.0):
    """Compute summary metrics."""
    loss, normalizer = cross_entropy(logits, labels, weights, label_smoothing)
    acc, _ = accuracy(logits, labels, weights)
    metrics = {"loss": loss, "accuracy": acc, "normalizer": normalizer}
    metrics = jax.lax.psum(metrics, axis_name="batch")
    return metrics