guides/distributed_training_with_jax.py/train_step
def train_step(train_state, x, y):
    (
        trainable_variables,
        non_trainable_variables,
        optimizer_variables,
    ) = train_state
    (loss_value, non_trainable_variables), grads = compute_gradients(
        trainable_variables, non_trainable_variables, x, y
    )

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

    return loss_value, (
        trainable_variables,
        non_trainable_variables,
        optimizer_variables,
    )