examples/demo_jax_distributed.py/train_step
def train_step(train_state, x, y):
    (loss_value, non_trainable_variables), grads = compute_gradients(
        train_state.trainable_variables,
        train_state.non_trainable_variables,
        x,
        y,
    )

    trainable_variables, optimizer_variables = optimizer.stateless_apply(
        train_state.optimizer_variables, grads, train_state.trainable_variables
    )

    return loss_value, TrainingState(
        trainable_variables, non_trainable_variables, optimizer_variables
    )