def eval_step(*, flax_model: nn.Module, train_state: train_utils.TrainState, batch: bert_base_model.Batch, metrics_fn: bert_base_model.MetricFn, all_gather: bool=False, debug: Optional[bool]=False) -> Tuple[(Dict[(str, Tuple[(float, int)])], Optional[jnp.ndarray], Optional[jnp.ndarray])]:
    'Runs a single step of training.\n\n  Note that in this code, the buffer of the second argument (batch) is donated\n  to the computation.\n\n  Assumed API of metrics_fn is:\n  ```metrics = metrics_fn(logits, batch)\n  where batch is yielded by the batch iterator, and metrics is a dictionary\n  mapping metric name to a vector of per example measurements. eval_step will\n  aggregate (by summing) all per example measurements and divide by the\n  aggregated normalizers. For each given metric we compute:\n  1/N sum_{b in batch_iter} metric(b), where  N is the sum of normalizer\n  over all batches.\n\n  Args:\n    flax_model: A Flax model.\n    train_state: TrainState, the state of training including the current\n      global_step, model_state, rng, and optimizer. The buffer of this argument\n      can be donated to the computation.\n    batch: A single batch of data. a metrics function, that given logits and\n      batch of data, calculates the metrics as well as 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    all_gather: If True, the function gather batch and output of\n      model in from all hosts, using `jax.lax.all_gather` and return it, e.g.,\n      for computing global metrics on CPU.\n    debug: Whether the debug mode is enabled during evaluation. `debug=True`\n      enables model specific logging/storing some values using\n      jax.host_callback.\n\n  Returns:\n    Calculated metrics and optionally output, and batch after all_gather.\n  '
    variables = {'params': train_state.optimizer.target, **train_state.model_state}
    output = flax_model.apply(variables, batch, train=False, mutable=False, debug=debug)
    metrics = metrics_fn(output, batch)
    if all_gather:
        output = jax.lax.all_gather(output, 'batch')
        batch = jax.lax.all_gather(batch, 'batch')
        return (metrics, output, batch)
    else:
        return (metrics, None, None)
Runs a single step of training.

Note that in this code, the buffer of the second argument (batch) is donated
to the computation.

Assumed API of metrics_fn is:
```metrics = metrics_fn(logits, batch)
where batch is yielded by the batch iterator, and metrics is a dictionary
mapping metric name to a vector of per example measurements. eval_step will
aggregate (by summing) all per example measurements and divide by the
aggregated normalizers. For each given metric we compute:
1/N sum_{b in batch_iter} metric(b), where  N is the sum of normalizer
over all batches.

Args:
  flax_model: A Flax model.
  train_state: TrainState, 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. a metrics function, that given logits and
    batch of data, calculates the metrics as well as the loss.
  metrics_fn: A metrics function, that given logits and batch of data,
    calculates the metrics as well as the loss.
  all_gather: If True, the function gather batch and output of
    model in from all hosts, using `jax.lax.all_gather` and return it, e.g.,
    for computing global metrics on CPU.
  debug: Whether the debug mode is enabled during evaluation. `debug=True`
    enables model specific logging/storing some values using
    jax.host_callback.

Returns:
  Calculated metrics and optionally output, and batch after all_gather.
def eval_step(*, flax_model: nn.Module, train_state: train_utils.TrainState, batch: bert_base_model.Batch, metrics_fn: bert_base_model.MetricFn, all_gather: bool=False, debug: Optional[bool]=False) -> Tuple[(Dict[(str, Tuple[(float, int)])], Optional[jnp.ndarray], Optional[jnp.ndarray])]:
    'Runs a single step of training.\n\n  Note that in this code, the buffer of the second argument (batch) is donated\n  to the computation.\n\n  Assumed API of metrics_fn is:\n  ```metrics = metrics_fn(logits, batch)\n  where batch is yielded by the batch iterator, and metrics is a dictionary\n  mapping metric name to a vector of per example measurements. eval_step will\n  aggregate (by summing) all per example measurements and divide by the\n  aggregated normalizers. For each given metric we compute:\n  1/N sum_{b in batch_iter} metric(b), where  N is the sum of normalizer\n  over all batches.\n\n  Args:\n    flax_model: A Flax model.\n    train_state: TrainState, the state of training including the current\n      global_step, model_state, rng, and optimizer. The buffer of this argument\n      can be donated to the computation.\n    batch: A single batch of data. a metrics function, that given logits and\n      batch of data, calculates the metrics as well as 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    all_gather: If True, the function gather batch and output of\n      model in from all hosts, using `jax.lax.all_gather` and return it, e.g.,\n      for computing global metrics on CPU.\n    debug: Whether the debug mode is enabled during evaluation. `debug=True`\n      enables model specific logging/storing some values using\n      jax.host_callback.\n\n  Returns:\n    Calculated metrics and optionally output, and batch after all_gather.\n  '
    variables = {'params': train_state.optimizer.target, **train_state.model_state}
    output = flax_model.apply(variables, batch, train=False, mutable=False, debug=debug)
    metrics = metrics_fn(output, batch)
    if all_gather:
        output = jax.lax.all_gather(output, 'batch')
        batch = jax.lax.all_gather(batch, 'batch')
        return (metrics, output, batch)
    else:
        return (metrics, None, None)