def batched_loss_fn(model):
    """Apply loss function across a batch of examples."""
    loss, metrics = jax.vmap(loss_fn, (None, 0, None))(model, batched_examples,
                                                       static_batch_metadata)
    return jnp.mean(loss), metrics