examples/research_projects/jax-projects/big_bird/bigbird_flax.py/loss_fn
def train_step(state, drp_rng, **model_inputs):
    def loss_fn(params):
        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=params, dropout_rng=drp_rng, train=True)
        start_logits, end_logits, pooled_logits = outputs

        return state.loss_fn(
            start_logits,
            start_labels,
            end_logits,
            end_labels,
            pooled_logits,
            pooled_labels,
        )

    drp_rng, new_drp_rng = jax.random.split(drp_rng)
    grad_fn = jax.value_and_grad(loss_fn)
    loss, grads = grad_fn(state.params)
    metrics = jax.lax.pmean({"loss": loss}, axis_name="batch")
    grads = jax.lax.pmean(grads, "batch")

    state = state.apply_gradients(grads=grads)
    return state, metrics, new_drp_rng