def train_step(*, flax_model: nn.Module, train_state: train_utils.TrainState, batch: bert_base_model.Batch, learning_rate_fn: Callable[([int], float)], loss_fn: bert_base_model.LossFn, metrics_fn: bert_base_model.MetricFn, config: ml_collections.ConfigDict, debug: Optional[bool]=False) -> Tuple[(train_utils.TrainState, Dict[(str, Tuple[(float, int)])], float)]:
    'Runs a single step of training.\n\n  Given the state of the training and a batch of data, computes\n  the loss and updates the parameters of the model.\n\n  Note that in this code, the buffers of the first (train_state) and second\n  (batch) arguments are donated to the computation.\n\n  Args:\n    flax_model: A Flax model.\n    train_state: The state of training including the current global_step,\n      model_state, rng, and optimizer. The buffer of this argument can be\n      donated to the computation.\n    batch: A single batch of data. The buffer of this argument can be donated to\n      the computation.\n    learning_rate_fn: Learning rate scheduler which given the global_step\n      generates the learning rate.\n    loss_fn: A loss function that given logits, a batch, and parameters of the\n      model calculates the loss.\n    metrics_fn: A metrics function that given logits and batch of data,\n      calculates the metrics as well as the loss.\n    config: Configurations of the experiment.\n    debug: Whether the debug mode is enabled during training. `debug=True`\n      enables model specific logging/storing some values using\n      jax.host_callback.\n\n  Returns:\n    Updated state of training, computed metrics, and learning rate for logging.\n  '
    (new_rng, rng) = jax.random.split(train_state.rng)
    dropout_rng = train_utils.bind_rng_to_host_device(rng, axis_name='batch', bind_to='device')

    def training_loss_fn(params):
        variables = {'params': params, **train_state.model_state}
        (output, new_model_state) = flax_model.apply(variables, batch, mutable=['batch_stats'], train=True, rngs={'dropout': dropout_rng}, debug=debug)
        loss = loss_fn(output, batch, variables['params'])
        return (loss, (new_model_state, output))
    compute_gradient_fn = jax.value_and_grad(training_loss_fn, has_aux=True)
    step = train_state.global_step
    lr = learning_rate_fn(step)
    ((train_cost, (new_model_state, output)), grad) = compute_gradient_fn(train_state.optimizer.target)
    del train_cost
    if (config.get('max_grad_norm', None) is not None):
        grad = jax_optimizers.clip_grads(grad, config.max_grad_norm)
    grad = jax.lax.pmean(grad, axis_name='batch')
    new_optimizer = train_state.optimizer.apply_gradient(grad, learning_rate=lr)
    if (config.get('explicit_weight_decay', None) is not None):
        new_optimizer = new_optimizer.replace(target=optimizers.tree_map_with_names(functools.partial(optimizers.decay_weight_fn, lr=lr, decay=config.explicit_weight_decay), new_optimizer.target, match_name_fn=(lambda name: ('kernel' in name))))
    metrics = metrics_fn(output, batch)
    new_train_state = train_state.replace(global_step=(step + 1), optimizer=new_optimizer, model_state=new_model_state, rng=new_rng)
    return (new_train_state, metrics, lr)
Runs a single step of training.

Given the state of the training and a batch of data, computes
the loss and updates the parameters of the model.

Note that in this code, the buffers of the first (train_state) and second
(batch) arguments are donated to the computation.

Args:
  flax_model: A Flax model.
  train_state: The state of training including the current global_step,
    model_state, rng, and optimizer. The buffer of this argument can be
    donated to the computation.
  batch: A single batch of data. The buffer of this argument can be donated to
    the computation.
  learning_rate_fn: Learning rate scheduler which given the global_step
    generates the learning rate.
  loss_fn: A loss function that given logits, a batch, and parameters of the
    model calculates the loss.
  metrics_fn: A metrics function that given logits and batch of data,
    calculates the metrics as well as the loss.
  config: Configurations of the experiment.
  debug: Whether the debug mode is enabled during training. `debug=True`
    enables model specific logging/storing some values using
    jax.host_callback.

Returns:
  Updated state of training, computed metrics, and learning rate for logging.
def train_step(*, flax_model: nn.Module, train_state: train_utils.TrainState, batch: bert_base_model.Batch, learning_rate_fn: Callable[([int], float)], loss_fn: bert_base_model.LossFn, metrics_fn: bert_base_model.MetricFn, config: ml_collections.ConfigDict, debug: Optional[bool]=False) -> Tuple[(train_utils.TrainState, Dict[(str, Tuple[(float, int)])], float)]:
    'Runs a single step of training.\n\n  Given the state of the training and a batch of data, computes\n  the loss and updates the parameters of the model.\n\n  Note that in this code, the buffers of the first (train_state) and second\n  (batch) arguments are donated to the computation.\n\n  Args:\n    flax_model: A Flax model.\n    train_state: The state of training including the current global_step,\n      model_state, rng, and optimizer. The buffer of this argument can be\n      donated to the computation.\n    batch: A single batch of data. The buffer of this argument can be donated to\n      the computation.\n    learning_rate_fn: Learning rate scheduler which given the global_step\n      generates the learning rate.\n    loss_fn: A loss function that given logits, a batch, and parameters of the\n      model calculates the loss.\n    metrics_fn: A metrics function that given logits and batch of data,\n      calculates the metrics as well as the loss.\n    config: Configurations of the experiment.\n    debug: Whether the debug mode is enabled during training. `debug=True`\n      enables model specific logging/storing some values using\n      jax.host_callback.\n\n  Returns:\n    Updated state of training, computed metrics, and learning rate for logging.\n  '
    (new_rng, rng) = jax.random.split(train_state.rng)
    dropout_rng = train_utils.bind_rng_to_host_device(rng, axis_name='batch', bind_to='device')

    def training_loss_fn(params):
        variables = {'params': params, **train_state.model_state}
        (output, new_model_state) = flax_model.apply(variables, batch, mutable=['batch_stats'], train=True, rngs={'dropout': dropout_rng}, debug=debug)
        loss = loss_fn(output, batch, variables['params'])
        return (loss, (new_model_state, output))
    compute_gradient_fn = jax.value_and_grad(training_loss_fn, has_aux=True)
    step = train_state.global_step
    lr = learning_rate_fn(step)
    ((train_cost, (new_model_state, output)), grad) = compute_gradient_fn(train_state.optimizer.target)
    del train_cost
    if (config.get('max_grad_norm', None) is not None):
        grad = jax_optimizers.clip_grads(grad, config.max_grad_norm)
    grad = jax.lax.pmean(grad, axis_name='batch')
    new_optimizer = train_state.optimizer.apply_gradient(grad, learning_rate=lr)
    if (config.get('explicit_weight_decay', None) is not None):
        new_optimizer = new_optimizer.replace(target=optimizers.tree_map_with_names(functools.partial(optimizers.decay_weight_fn, lr=lr, decay=config.explicit_weight_decay), new_optimizer.target, match_name_fn=(lambda name: ('kernel' in name))))
    metrics = metrics_fn(output, batch)
    new_train_state = train_state.replace(global_step=(step + 1), optimizer=new_optimizer, model_state=new_model_state, rng=new_rng)
    return (new_train_state, metrics, lr)