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 build_hparams(model_name,
                  initializer_name,
                  dataset_name,
                  hparam_file,
                  hparam_overrides): 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): _ALL_DATASETS = {
    'mnist':
        _Dataset(small_image_datasets.get_mnist,
                 small_image_datasets.MNIST_HPARAMS,
                 small_image_datasets.MNIST_METADATA),
    'mnist_autoencoder':
        _Dataset(small_image_datasets.get_mnist_autoencoder,
                 small_image_datasets.MNIST_AUTOENCODER_HPARAMS,
                 small_image_datasets.MNIST_AUTOENCODER_METADATA),
    'fashion_mnist':
        _Dataset(small_image_datasets.get_fashion_mnist,
                 small_image_datasets.FASHION_MNIST_HPARAMS,
                 small_image_datasets.FASHION_MNIST_METADATA),
    'cifar10':
        _Dataset(small_image_datasets.get_cifar10,
                 small_image_datasets.CIFAR10_DEFAULT_HPARAMS,
                 small_image_datasets.CIFAR10_METADATA),
    'cifar100':
        _Dataset(small_image_datasets.get_cifar100,
                 small_image_datasets.CIFAR100_DEFAULT_HPARAMS,
                 small_image_datasets.CIFAR100_METADATA),
    'criteo1tb':
        _Dataset(criteo_terabyte_input_pipeline.get_criteo1tb,
                 criteo_terabyte_input_pipeline.CRITEO1TB_DEFAULT_HPARAMS,
                 criteo_terabyte_input_pipeline.CRITEO1TB_METADATA),
    'fake':
        _Dataset(fake_dataset.get_fake, fake_dataset.DEFAULT_HPARAMS,
                 fake_dataset.METADATA),
    'imagenet':
        _Dataset(imagenet_dataset.get_imagenet,
                 imagenet_dataset.DEFAULT_HPARAMS, imagenet_dataset.METADATA),
    'translate_wmt':
        _Dataset(translate_wmt.get_translate_wmt, translate_wmt.DEFAULT_HPARAMS,
                 translate_wmt.METADATA),
    'mlperf_imagenet':
        _Dataset(mlperf_imagenet_dataset.get_mlperf_imagenet,
                 mlperf_imagenet_dataset.DEFAULT_HPARAMS,
                 mlperf_imagenet_dataset.METADATA),
    'svhn_no_extra':
        _Dataset(small_image_datasets.get_svhn_no_extra,
                 small_image_datasets.SVHN_NO_EXTRA_DEFAULT_HPARAMS,
                 small_image_datasets.SVHN_NO_EXTRA_METADATA),
    'nqm_noise':
        _Dataset(nqm_noise.get_nqm_noise, nqm_noise.NQM_HPARAMS,
                 nqm_noise.NQM_METADATA),
    'ogbg_molpcba':
        _Dataset(ogbg_molpcba.get_ogbg_molpcba, ogbg_molpcba.DEFAULT_HPARAMS,
                 ogbg_molpcba.METADATA),
    'uniref50':
        _Dataset(proteins.get_uniref, proteins.DEFAULT_HPARAMS,
                 proteins.METADATA),
}
def get_dataset(dataset_name):
def get_dataset_hparams(dataset_name):
def get_dataset_meta_data(dataset_name): MNIST_HPARAMS = config_dict.ConfigDict(dict(
    train_size=50000,
    valid_size=10000,
    test_size=10000,
    input_shape=(28, 28, 1),
    output_shape=(10,)))
MNIST_METADATA = {
    'apply_one_hot_in_loss': False,
}
MNIST_AUTOENCODER_HPARAMS = config_dict.ConfigDict(dict(
    train_size=50000,
    valid_size=10000,
    test_size=10000,
    input_shape=(28, 28, 1),
    output_shape=(784,)))
MNIST_AUTOENCODER_METADATA = {
    'apply_one_hot_in_loss': False,
}
FASHION_MNIST_HPARAMS = config_dict.ConfigDict(dict(
    train_size=45000,
    valid_size=5000,
    test_size=10000,
    input_shape=(28, 28, 1),
    output_shape=(10,)))
FASHION_MNIST_METADATA = {
    'apply_one_hot_in_loss': False,
}
CIFAR10_DEFAULT_HPARAMS = config_dict.ConfigDict(dict(
    flip_probability=0.5,
    alpha=1.0,
    crop_num_pixels=4,
    use_mixup=True,
    train_size=45000,
    valid_size=5000,
    test_size=10000,
    input_shape=(32, 32, 3),
    output_shape=(10,)))
CIFAR10_METADATA = {
    'apply_one_hot_in_loss': False,
}
CIFAR100_DEFAULT_HPARAMS = config_dict.ConfigDict(dict(
    flip_probability=0.5,
    alpha=1.0,
    crop_num_pixels=4,
    use_mixup=True,
    train_size=45000,
    valid_size=5000,
    test_size=10000,
    input_shape=(32, 32, 3),
    output_shape=(100,)))
CIFAR100_METADATA = {
    'apply_one_hot_in_loss': False,
}
SVHN_NO_EXTRA_DEFAULT_HPARAMS = config_dict.ConfigDict(dict(
    flip_probability=0.5,
    alpha=1.0,
    crop_num_pixels=4,
    use_mixup=False,
    train_size=73257 - 7000,
    valid_size=7000,
    test_size=26032,
    input_shape=(32, 32, 3),
    output_shape=(10,)))
SVHN_NO_EXTRA_METADATA = {
    'apply_one_hot_in_loss': False,
}
def _eval_batches(images,
                  labels,
                  per_host_batch_size,
                  num_batches=None,
                  valid_example_keys=None):
def _shard_by_host_id(array):
def _prepare_small_image_datasets(
    data_train,
    data_valid,
    data_test,
    per_host_batch_size,
    per_host_eval_batch_size,
    train_size,
    rescale,
    input_shape,
    output_shape,
    shuffle_rng,
    augment_fn,
    is_one_hot=True,
    autoencoder=False,
    include_example_keys=False):
def _process_small_tfds_image_ds(
    dataset_name,
    per_host_batch_size,
    per_host_eval_batch_size,
    train_size,
    valid_size,
    test_size,
    rescale,
    input_shape,
    output_shape,
    shuffle_rng,
    augment_fn=image_preprocessing.identity_augment,
    is_one_hot=True,
    autoencoder=False):
def _load_deterministic_with_custom_validation(
    dataset_name,
    per_host_batch_size,
    per_host_eval_batch_size,
    train_size,
    valid_size,
    test_size,
    rescale,
    input_shape,
    output_shape,
    shuffle_rng,
    augment_fn=image_preprocessing.identity_augment,
    include_example_keys=False):
def get_mnist(shuffle_rng, batch_size, eval_batch_size, hps=None):
def get_mnist_autoencoder(shuffle_rng, batch_size, eval_batch_size, hps=None):
def get_fashion_mnist(shuffle_rng,
                      batch_size,
                      eval_batch_size,
                      hps=None):
def get_cifar10(shuffle_rng,
                batch_size,
                eval_batch_size,
                hps=None):
def get_cifar100(shuffle_rng,
                 batch_size,
                 eval_batch_size,
                 hps=None):
def get_svhn_no_extra(
    shuffle_rng,
    batch_size,
    eval_batch_size,
    hps=None): def initialize(flax_module,
               initializer,
               loss_fn,
               input_shape,
               output_shape,
               hps,
               rng,
               metrics_logger,
               fake_batch=None): def noop(
    loss_fn=None,
    flax_module=None,
    params=None,
    hps=None,
    input_shape=None,
    output_shape=None,
    rng_key=None,
    metrics_logger=None,
):
def get_initializer(initializer_name):
def get_initializer_hparams(initializer_name):
DEFAULT_HPARAMS = config_dict.ConfigDict()
_ALL_INITIALIZERS = {
    'noop': (noop, DEFAULT_HPARAMS),
    'meta_init': (meta_init.meta_init, meta_init.DEFAULT_HPARAMS),
    'sparse_init': (sparse_init.sparse_init, sparse_init.DEFAULT_HPARAMS),
} def _evaluate_batch(flax_module, params, batch_stats, batch, metrics_bundle,
                    apply_one_hot_in_loss):
def _predict_batch(flax_module,
                   params,
                   batch_stats,
                   batch,
                   output_activation_fn=None):
  def __init__(self, hps, dataset_meta_data, loss_name, metrics_name):
  def evaluate_batch(self, params, batch_stats, batch):
  def predict_batch(self,
                    params,
                    batch_stats,
                    batch,
                    apply_output_activation_fn=False):
  def training_cost(self, params, batch, batch_stats=None, dropout_rng=None):
  def get_fake_batch(self, hps):
  def build_flax_module(self):
class BaseModel(object): class NumExamples(metrics.Metric):
  class _Metric(metrics.Metric):
  def from_model_output(cls, logits, targets, weights, **_):
  def merge(self, other):
  def compute(self):
def weighted_average_metric(fun):
    def from_model_output(cls, logits, targets, weights, **_):
    def merge(self, other):
    def compute(self):
def weighted_misclassifications(logits, targets, weights=None):
def get_metrics(metrics_name):
_METRICS = {
    'binary_autoencoder_metrics':
        metrics.Collection.create(
            sigmoid_binary_cross_entropy=weighted_average_metric(
                losses.unnormalized_sigmoid_binary_cross_entropy),
            sigmoid_mean_squared_error=weighted_average_metric(
                losses.unnormalized_sigmoid_mean_squared_error),
            num_examples=NumExamples),
    'classification_metrics':
        metrics.Collection.create(
            error_rate=weighted_average_metric(weighted_misclassifications),
            ce_loss=weighted_average_metric(
                losses.weighted_unnormalized_cross_entropy),
            num_examples=NumExamples),
    'binary_classification_metrics':
        metrics.Collection.create(
            ce_loss=weighted_average_metric(
                losses.sigmoid_binary_cross_entropy),
            num_examples=NumExamples,
            average_precision=losses.MeanAveragePrecision)
} _ALL_MODELS = {
    'fully_connected':
        (fully_connected.FullyConnectedModel, fully_connected.DEFAULT_HPARAMS),
    'simple_cnn': (simple_cnn.SimpleCNNModel, simple_cnn.DEFAULT_HPARAMS),
    'max_pooling_cnn':
        (max_pooling_cnn.MaxPoolingCNNModel, max_pooling_cnn.DEFAULT_HPARAMS),
    'wide_resnet': (wide_resnet.WideResnetModel, wide_resnet.DEFAULT_HPARAMS),
    'resnet': (resnet.ResnetModel, resnet.DEFAULT_HPARAMS),
    'adabelief_densenet': (adabelief_densenet.AdaBeliefDensenetModel,
                           adabelief_densenet.DEFAULT_HPARAMS),
    'adabelief_resnet': (adabelief_resnet.AdaBeliefResnetModel,
                         adabelief_resnet.DEFAULT_HPARAMS),
    'adabelief_vgg':
        (adabelief_vgg.AdaBeliefVGGModel, adabelief_vgg.DEFAULT_HPARAMS),
    'autoencoder': (autoencoder.AutoEncoderModel, autoencoder.DEFAULT_HPARAMS),
    'convolutional_autoencoder':
        (convolutional_autoencoder.ConvAutoEncoderModel,
         convolutional_autoencoder.DEFAULT_HPARAMS),
    'fake_resnet':
        (mlperf_resnet.FakeModel, mlperf_resnet.FAKE_MODEL_DEFAULT_HPARAMS),
    'mlperf_resnet':
        (mlperf_resnet.ResnetModelMLPerf, mlperf_resnet.MLPERF_DEFAULT_HPARAMS),
    'transformer':
        (transformer_lm.TransformerLM1B, transformer_lm.DEFAULT_HPARAMS),
    'nqm': (nqm.NQM, nqm.DEFAULT_HPARAMS),
    'xformer_translate': (xformer_translate.TransformerTranslate,
                          xformer_translate.DEFAULT_HPARAMS),
    'gnn': (gnn.GNNModel, gnn.DEFAULT_HPARAMS),
    'dlrm': (dlrm.DLRMModel, dlrm.DEFAULT_HPARAMS),
    'vit': (vit.ViTModel, vit.DEFAULT_HPARAMS),
}
def get_model(model_name):
def get_model_hparams(model_name): _GRAD_CLIP_EPS = 1e-6
def evaluate(
    params,
    batch_stats,
    batch_iter,
    evaluate_batch_pmapped):
def _inject_learning_rate(optimizer_state, lr):
def update(
    optimizer_state,
    params,
    batch_stats,
    batch,
    step,
    lr,
    rng,
    local_device_index,
    training_metrics_grabber,
    training_cost,
    grad_clip,
    optimizer_update_fn):
  def opt_cost(params):
def _merge_and_apply_prefix(d1, d2, prefix):
def eval_metrics(params, batch_stats, dataset, eval_num_batches,
                 eval_train_num_batches, evaluate_batch_pmapped):
def train(train_dir,
          model,
          dataset_builder,
          initializer,
          num_train_steps,
          hps,
          rng,
          eval_batch_size,
          eval_num_batches,
          eval_train_num_batches,
          eval_frequency,
          checkpoint_steps,
          early_stopping_target_name=None,
          early_stopping_target_value=None,
          early_stopping_mode=None,
          eval_steps=None,
          metrics_logger=None,
          init_logger=None,
          training_metrics_config=None,
          callback_configs=None):
  def get_step_frequency(cur_step):
import copy
import functools
import itertools
import os
import shutil
import tempfile
import jax.numpy as jnp
import jax.random
import jraph
import numpy as np
import pandas
import tensorflow.compat.v1 as tf  # importing this is needed for tfds mocking.
import tensorflow_datasets as tfds
from absl import flags
from absl.testing import absltest
from absl.testing import parameterized
from flax import jax_utils
from flax import linen as nn
from init2winit import checkpoint
from init2winit import hyperparameters
from init2winit import utils
from init2winit.dataset_lib import datasets
from init2winit.dataset_lib.small_image_datasets import Dataset
from init2winit.init_lib import init_utils
from init2winit.init_lib import initializers
from init2winit.model_lib import base_model
from init2winit.model_lib import metrics
from init2winit.model_lib import models
from init2winit.trainer_lib import trainer
from ml_collections.config_dict import config_dict
  meta_data = {
      'apply_one_hot_in_loss': True,
      'shift_inputs': True,
      'causal': True
  }
  return (Dataset(train_iterator_fn, eval_train_epoch, valid_epoch,
                  test_epoch), meta_data)


def _get_fake_graph_dataset(batch_size, eval_num_batches, hps):
  """Yields graph batches with label 1 if graph has 4 or 6 nodes and 0 otherwise."""

  def _get_batch(n_nodes_list, batch_size):
    # Hardcode batch_size to be 2 for simplicity
    del batch_size
    graphs_list = []
    labels_list = []
    weights_list = []
    for n_nodes in n_nodes_list:
      n_edges = n_nodes**2
      graph = jraph.get_fully_connected_graph(
          n_nodes, 1,
          np.ones((n_nodes, *hps.input_node_shape)))
      graph = graph._replace(
          edges=np.ones((n_edges, *hps.input_edge_shape)))
      labels = np.ones(hps.output_shape) * (1 if n_nodes in [4, 6] else 0)
      weights = np.ones(*hps.output_shape)
      graphs_list.append(graph)
      labels_list.append(labels)
      weights_list.append(weights)
    return {
        'inputs': jraph.batch(graphs_list),
        'targets': np.stack(labels_list),
        'weights': np.stack(weights_list),
    }

  # Ensure each batch has one positive and one negative example for
  # average precision to not be NaN.
  def train_iterator_fn():
    for ns in itertools.cycle([[3, 6], [4, 5]]):
      yield _get_batch(ns, batch_size)

  def eval_train_epoch(num_batches=None):
    if num_batches is None:
      num_batches = eval_num_batches
    for ns in itertools.islice(itertools.cycle([[3, 6], [4, 5]]), num_batches):
      yield _get_batch(ns, batch_size)

  def valid_epoch(num_batches=None):
    if num_batches is None:
      num_batches = eval_num_batches
    for ns in itertools.islice(itertools.cycle([[3, 6], [4, 5]]), num_batches):
      yield _get_batch(ns, batch_size)

  def test_epoch(num_batches=None):
    if num_batches is None:
      num_batches = eval_num_batches
    for ns in itertools.islice(itertools.cycle([[3, 6], [4, 5]]), num_batches):
      yield _get_batch(ns, batch_size)

  return (Dataset(train_iterator_fn, eval_train_epoch, valid_epoch,
                  test_epoch), {
                      'apply_one_hot_in_loss': False,
                  })


def _get_fake_dlrm_dataset(batch_size, eval_num_batches, hps):
  """Yields a single text batch repeatedly for train and test."""
  cat_features = []
  for vocab_size in hps.vocab_sizes:
    cat_features.append(
        np.random.randint(low=0, high=vocab_size, size=(batch_size, 1)))
  cat_features = np.concatenate(cat_features, 1)
  int_features = np.random.normal(size=(batch_size, hps.num_dense_features))
  inputs = np.concatenate((int_features, cat_features), 1)
  targets = np.random.randint(low=0, high=2, size=(batch_size, 1))
  batch = {
      'inputs': inputs,
      'targets': targets,
      'weights': jnp.ones(targets.shape),
  }

  def train_iterator_fn():
    while True:
      yield batch

  def eval_train_epoch(num_batches=None):
    if num_batches is None:
      num_batches = eval_num_batches
    for _ in range(num_batches):
      yield batch

  def valid_epoch(num_batches=None):
    if num_batches is None:
      num_batches = eval_num_batches
    for _ in range(num_batches):
      yield batch

  def test_epoch(num_batches=None):
    if num_batches is None:
      num_batches = eval_num_batches
    for _ in range(num_batches):
      yield batch

  meta_data = {
      'apply_one_hot_in_loss': False,
  }
  return (Dataset(train_iterator_fn, eval_train_epoch, valid_epoch,
                  test_epoch), meta_data)


class TrainerTest(parameterized.TestCase):
  """Tests training for 2 epochs on MNIST."""

  def setUp(self):
    super(TrainerTest, self).setUp()
    self.test_dir = tempfile.mkdtemp()

  def tearDown(self):
    # To not delete a directory in which a file might be created:

    checkpoint.wait_for_checkpoint_save()