def parallel_metrics_batch(model, batched_examples, batch_mask,
                             static_metadata):
    loss, metrics = jax.vmap(loss_fn, (None, 0, None))(model, batched_examples,
                                                       static_metadata)
    metrics["loss"] = loss
    metrics = jax.tree_map(
        lambda x: jnp.where(batch_mask, x, jnp.zeros_like(x)), metrics)
    metrics = jax.tree_map(lambda x: jax.lax.psum(jnp.sum(x), "devices"),
                           metrics)
    return metrics