_ALL_CALLBACKS = {
    'hessian': hessian_callback.HessianCallback,
    'mt': mt_callback.MTEvaluationCallback,
    'model_debugger': model_debugger_callback.ModelDebugCallback,
}
def get_callback(callback_name): FLAGS = flags.FLAGS
def replicate_checkpoint(
    latest,
    pytree_keys: Sequence[str],
    replicate=True):
def replicate_and_maybe_restore_latest_checkpoint(
    unreplicated_optimizer_state,
    unreplicated_params,
    unreplicated_batch_stats,
    unreplicated_training_metrics_grabber,
    train_dir):
def save_unreplicated_checkpoint_background(
    train_dir,
    optimizer_state,
    params,
    batch_stats,
    training_metrics_grabber,
    global_step,
    preemption_count,
    sum_train_cost,
    max_to_keep=1):
def _save_checkpoint_background_catch_error(*args, **kwargs):
def wait_for_checkpoint_save():
def save_checkpoint_background(*args, **kwargs):
def save_checkpoint(train_dir,
                    step,
                    state,
                    prefix='ckpt_',
                    max_to_keep=None):
def load_checkpoint(
    checkpoint_path,
    target=None,
    prefix='ckpt_'):
def load_latest_checkpoint(train_dir, target=None, prefix='ckpt_'): def _check_schedule_hparams(schedule_hparams, expected_keys):
def constant_schedule(schedule_hparams, max_training_steps):
def cosine_schedule(schedule_hparams, max_training_steps):
  def lr_fn(t):
def rsqrt_normalized_decay(schedule_hparams, max_training_steps):
  def lr_fn(t):
def polynomial_schedule(schedule_hparams, max_training_steps):
  def lr_fn(t):
def piecewise_constant_schedule(schedule_hparams, max_training_steps):
  def lr_fn(t):
def piecewise_linear_schedule(schedule_hparams, max_training_steps):
  def lr_fn(t):
def mlperf_polynomial_schedule(schedule_hparams, max_training_steps):
  def step_fn(step):
def compound_schedule(schedule_hparams, max_training_steps):
  def lr_fn(step):
def prepend_linear_warmup(schedule_hparams, max_training_steps,
                          base_lr_schedule):
  def lr_fn(t):
def warmup_then_piecewise_constant_schedule(schedule_hparams,
                                            max_training_steps):
def schedule_stretcher(schedule_fn, stretch_factor):
  def _schedule(global_step):
def get_schedule_fn(schedule_hparams, max_training_steps, stretch_factor=1): class TrainingDivergedError(Exception):
class _MetricsLeafState:
class TrainingMetricsGrabber:
class MetricLogger(object):
def dtype_from_str(dtype_string):
def timed(f):
  def wrapper(*args, **kwargs):
def set_up_loggers(train_dir, xm_work_unit=None):
def get_summary_tree(training_metrics_grabber):
  def _reduce_node(node):
def _update_param_stats(leaf_state, param_gradient, update, new_param, config):
def _validate_config(config):
  def create(model_params, config):
    def _node_init(x):
  def update(self, model_gradient, old_params, new_params):
  def state_dict(self):
  def restore_state(self, state, state_dict):
def run_in_parallel(function, list_of_kwargs_to_function, num_workers):
def should_eval(global_step, eval_frequency, eval_steps):
def add_log_file(logfile):
def load_pytrees(pytree_path):
  def __init__(self,
               csv_path='',
               json_path='',
               pytree_path='',
               events_dir=None,
               **logger_kwargs):
  def append_scalar_metrics(self, metrics):
  def write_pytree(self, pytree):
  def append_pytree(self, pytree):
  def append_json_object(self, json_obj):
def _summary_str(param):
def log_pytree_shape_and_statistics(pytree, json_path=None): def iterator_as_numpy(iterator):
def image_iterator(data,
                   rescale,
                   output_shape,
                   is_one_hot,
                   autoencoder,
                   shuffle_rng=None,
                   augment_fn=None,
                   include_example_keys=False):
def maybe_pad_batch(batch,
                    desired_batch_size,
                    data_format=None,
                    mask_key=None):
  def zero_pad(ar, pad_axis):
def shard(batch, n_devices=None):
  def _shard_array(array):
def tf_to_numpy(tfds_data): def initialize(flax_module,
               initializer,
               loss_fn,
               input_shape,
               output_shape,
               hps,
               rng,
               metrics_logger,
               fake_batch=None): ACTIVATIONS = {
    'relu': jax.nn.relu,
    'id': lambda x: x,
    'tanh': jnp.tanh,
    'sigmoid': jax.nn.sigmoid,
    'silu': jax.nn.silu,
    'leaky_relu': jax.nn.leaky_relu,
}
INITIALIZERS = {
    'delta_orthogonal': initializers.delta_orthogonal,
    'orthogonal': initializers.orthogonal,
    'lecun_normal': lecun_normal,
    'xavier_uniform': initializers.xavier_uniform,
}
class ScalarMultiply(nn.Module):
  def __call__(self, x):
def get_normalizer(normalizer,
                   train,
                   batch_size=None,
                   virtual_batch_size=None,
                   total_batch_size=None,
                   dtype=jnp.float32,
                   data_format='NHWC'):
    def identity_wrapper(*args, **kwargs):
      def identity(x, *args, **kwargs):
def apply_label_smoothing(one_hot_targets, label_smoothing):
def sync_batchnorm_stats(state):
def cross_device_avg(pytree):
def l2_regularization(params, l2_decay_rank_threshold):
def flatten_dict(nested_dict, sep='/'):
def rescale_layers(params, layer_rescale_factors): class GradientAccumulatorState(NamedTuple):
def accumulate_gradients(
    per_step_batch_size: int,
    total_batch_size: int,
    virtual_batch_size: Optional[int],
    base_opt_init_fn: optax.TransformInitFn,
    base_opt_update_fn: optax.TransformUpdateFn
) -> optax.GradientTransformation:
  def init_fn(params):
  def update_fn(updates, state, params=None):
    def total_batch_update(total_gradients, params, state):
    def accumulation_continuation(updated_accumulations, _, state): def sgd(learning_rate, weight_decay, momentum=None, nesterov=False):
def get_optimizer(hps, model=None):
def _wrap_update_fn(opt_name, opt_update):
  def update_fn(grads, optimizer_state, params, batch=None, batch_stats=None): def log_epoch_report(report, metrics_logger):
def maybe_log_training_metrics(training_metrics_grabber, metrics_logger):
def maybe_sync_batchnorm_stats(batch_stats):
def should_eval(global_step, eval_frequency, eval_steps):
def check_for_early_stopping(
    early_stopping_target_name,
    early_stopping_target_value,
    early_stopping_mode,
    eval_report):
import functools
import itertools
import os
import time
import jax
import jax.numpy as jnp
import numpy as np
import optax
from absl import logging
from init2winit import callbacks
from init2winit import checkpoint
from init2winit import schedules
from init2winit import utils
from init2winit.dataset_lib import data_utils
from init2winit.init_lib import init_utils
from init2winit.model_lib import model_utils
from init2winit.optimizer_lib import gradient_accumulator
from init2winit.optimizer_lib import optimizers
from init2winit.trainer_lib import trainer_utils
from jax import lax
      grad_clip=hps.get('grad_clip'),
      optimizer_update_fn=optimizer_update_fn)
  # in_axes = (
  #     optimizer_state = 0,
  #     params = 0,
  #     batch_stats = 0,
  #     batch = 0,
  #     step = None,
  #     lr = None,
  #     rng = None,
  #     local_device_index = 0,
  #     training_metrics_grabber = 0,
  #     training_cost,
  #     grad_clip,
  #     optimizer_update_fn)
  # Also, we can donate buffers for 'optimizer', 'batch_stats',
  # 'batch' and 'training_metrics_grabber' for update's pmapped computation.
  update_pmapped = jax.pmap(
      update_fn,
      axis_name='batch',
      in_axes=(0, 0, 0, 0, None, None, None, 0, 0),
      donate_argnums=(0, 1, 2, 7))
  # During eval, we can donate the 'batch' buffer. We don't donate the
  # 'params' and 'batch_stats' buffers as we don't re-assign those values in
  # eval, we do that only in train.
  evaluate_batch_pmapped = jax.pmap(
      model.evaluate_batch, axis_name='batch', donate_argnums=(2,))
  start_time = time.time()
  start_step = global_step
  prev_eval_step = start_step
  def get_step_frequency(cur_step):
    return float(cur_step - start_step) / (time.time() - start_time)

  if jax.process_index() == 0:
    logging.info('Starting training!')

  # Numpy array of range(0, local_device_count) to send to each device to be
  # folded into the RNG inside each train step to get a unique per-device RNG.
  local_device_indices = np.arange(jax.local_device_count())

  # Start at the resumed step and continue until we have finished the number of
  # training steps. If building a dataset iterator using a tf.data.Dataset, in
  # the case of a batch size that does not evenly divide the training dataset
  # size, if using `ds.batch(..., drop_remainer=True)` on the training dataset
  # then the final batch in this iterator will be a partial batch. However, if
  # `drop_remainer=False`, then this iterator will always return batches of the
  # same size, and the final batch will have elements from the start of the
  # (num_epochs + 1)-th epoch.
  train_iter = itertools.islice(
      dataset.train_iterator_fn(), global_step, num_train_steps)

  eval_callbacks = []
  rng, callback_rng = jax.random.split(rng)
  callback_rngs = jax.random.split(callback_rng, len(callback_configs))
  for callback_rng, config in zip(callback_rngs, callback_configs):
    eval_callback = callbacks.get_callback(
        config['callback_name'])(model, params, batch_stats,
                                 dataset, hps, config, train_dir, callback_rng)
    eval_callbacks.append(eval_callback)

  for batch in train_iter:
    if global_step in checkpoint_steps and jax.process_index() == 0:
      checkpoint.save_unreplicated_checkpoint_background(
          checkpoint_dir,
          optimizer_state,
          params,
          batch_stats,
          training_metrics_grabber,
          global_step,
          preemption_count,
          sum_train_cost,
          max_to_keep=None)
    batch = data_utils.shard(batch)
    lr = lr_fn(global_step)
    optimizer_state, params, batch_stats, cost_val, training_metrics_grabber, grad_norm = update_pmapped(
        optimizer_state,
        params,
        batch_stats,
        batch,
        global_step,
        lr,
        rng,
        local_device_indices,
        training_metrics_grabber)
    # Calling float is needed since cost_val is a shape (1,) DeviceArray.
    sum_train_cost += float(np.mean(cost_val))
    global_step += 1
    # TODO(gdahl, gilmer): consider moving this test up.
    # NB: Since this test is after we increment global_step, having 0 in
    # eval_steps does nothing.
    if trainer_utils.should_eval(global_step, eval_frequency, eval_steps):
      batch_stats = trainer_utils.maybe_sync_batchnorm_stats(batch_stats)
      report, eval_time = eval_metrics(params,
                                       batch_stats,
                                       dataset,
                                       eval_num_batches,
                                       eval_train_num_batches,
                                       evaluate_batch_pmapped)
      mean_train_cost = sum_train_cost / max(1, global_step - prev_eval_step)
      report.update(learning_rate=float(lr),
                    global_step=global_step,
                    epoch=global_step * hps.batch_size // hps.train_size,
                    steps_per_sec=get_step_frequency(global_step),
                    eval_time=eval_time,
                    grad_norm=np.mean(grad_norm),
                    preemption_count=preemption_count,
                    train_cost=mean_train_cost)

      for eval_callback in eval_callbacks:
        callback_metrics = eval_callback.run_eval(params, batch_stats,
                                                  global_step)
        if set(callback_metrics.keys()).intersection(set(report.keys())):
          raise ValueError('There was a collision between the callback metrics'
                           'and the standard eval metrics keys')
        report.update(callback_metrics)
      yield report
      if jax.process_index() == 0:
        trainer_utils.log_epoch_report(report, metrics_logger)
        trainer_utils.maybe_log_training_metrics(
            training_metrics_grabber, metrics_logger)

        checkpoint.save_unreplicated_checkpoint_background(