examples/research_projects/jax-projects/big_bird/bigbird_flax.py/val_step
def val_step(state, **model_inputs):
    start_labels = model_inputs.pop("start_labels")
    end_labels = model_inputs.pop("end_labels")
    pooled_labels = model_inputs.pop("pooled_labels")

    outputs = state.apply_fn(**model_inputs, params=state.params, train=False)
    start_logits, end_logits, pooled_logits = outputs

    loss = state.loss_fn(start_logits, start_labels, end_logits, end_labels, pooled_logits, pooled_labels)
    metrics = jax.lax.pmean({"loss": loss}, axis_name="batch")
    return metrics