examples/research_projects/performer/run_mlm_performer.py/eval_step
def eval_step(params, batch):
    """
    Calculate evaluation metrics on a batch.
    """
    targets = batch.pop("labels")

    # Hide away tokens which doesn't participate in the optimization
    token_mask = jnp.where(targets > 0, 1.0, 0.0)
    logits = model(**batch, params=params, train=False)[0]

    return compute_metrics(logits, targets, token_mask)