examples/research_projects/performer/run_mlm_performer.py/training_step
def training_step(optimizer, batch, dropout_rng):
    dropout_rng, new_dropout_rng = jax.random.split(dropout_rng)

    def loss_fn(params):
        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, dropout_rng=dropout_rng, train=True)[0]
        loss, weight_sum = cross_entropy(logits, targets, token_mask)
        return loss / weight_sum

    step = optimizer.state.step
    lr = lr_scheduler_fn(step)
    grad_fn = jax.value_and_grad(loss_fn)
    loss, grad = grad_fn(optimizer.target)
    grad = jax.lax.pmean(grad, "batch")
    optimizer = optimizer.apply_gradient(grad, learning_rate=lr)

    return loss, optimizer, new_dropout_rng