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
    diff = np.linalg.norm(v1.reshape(-1) - v2.reshape(-1))
    self.assertAlmostEqual(diff, 0.0)

  def test_classifaction_model_evaluate(self):
    """Test trainer evaluate end to end with classification model metrics."""
    # Define a fake model that always outputs the same logits.
    fake_batch_logits = np.tile([.5, .2, .7, 0.0], (4, 1))

    class FakeModel(nn.Module):

      @nn.compact
      def __call__(self, x, train):
        # Make a single linear layer with the identity as the init.
        identity_fn = lambda *_: np.eye(4)
        x = nn.Dense(features=4, use_bias=False, kernel_init=identity_fn)(
            fake_batch_logits)
        return x

    key = jax.random.PRNGKey(0)
    params_rng, dropout_rng = jax.random.split(key, num=2)
    fake_flax_module = FakeModel()
    model_init_fn = jax.jit(fake_flax_module.init)
    init_dict = model_init_fn(
        rngs={'params': params_rng, 'dropout': dropout_rng},
        x=None,
        train=False)
    params = jax_utils.replicate(init_dict['params'])
    batch_stats = init_dict.get('batch_stats', {})

    # 4 evaluation batches of size 4.
    weights = np.ones((4))
    fake_batches = [
        {
            'inputs': None,
            'targets': np.array([3, 2, 1, 0]),
            'weights': weights
        },
        {
            'inputs': None,
            'targets': np.array([0, 3, 2, 0]),
            'weights': weights
        },
        {
            'inputs': None,
            'targets': np.array([0, 0, 0, 0]),
            'weights': weights
        },
        {
            'inputs': None,
            'targets': np.array([1, 1, 1, 1]),
            'weights': weights
        },
    ]

    def fake_batches_gen():
      for batch in fake_batches:
        yield batch

    # pylint: disable=protected-access
    eval_fn = functools.partial(
        base_model._evaluate_batch,
        flax_module=fake_flax_module,
        metrics_bundle=metrics.get_metrics('classification_metrics'),
        apply_one_hot_in_loss=True)
    evaluate_batch_pmapped = jax.pmap(eval_fn, axis_name='batch')
    # pylint: enable=protected-access
    evaluated_metrics = trainer.evaluate(params,
                                         batch_stats,
                                         fake_batches_gen(),
                                         evaluate_batch_pmapped)

    def batch_ce_loss(logits, targets):
      one_hot_targets = np.eye(4)[targets]
      loss = -np.sum(one_hot_targets * nn.log_softmax(logits), axis=-1)
      return loss

    expected_error_rate = 14.0/16.0  # FakeModel always predicts class 2.
    expected_ce_loss = np.mean(
        [batch_ce_loss(fake_batch_logits, b['targets']) for b in fake_batches])

    self.assertEqual(expected_error_rate, evaluated_metrics['error_rate'])
    self.assertAlmostEqual(
        expected_ce_loss, evaluated_metrics['ce_loss'], places=4)
    self.assertEqual(16, evaluated_metrics['num_examples'])

  def test_graph_model_trainer(self):
    """Tests that graph model training decreases loss."""
    rng = jax.random.PRNGKey(1337)
    model_str = 'gnn'
    model_cls = models.get_model(model_str)
    hps = models.get_model_hparams(model_str)
    hps.update({
        'batch_size': 2,
        'input_edge_shape': (7,),
        'input_node_shape': (3,),
        'input_shape': (7, 3),
        'output_shape': (5,),
        'model_dtype': 'float32',
        'train_size': 15,
        'valid_size': 10,
        'test_size': 10,
        'num_message_passing_steps': 1,
        'normalizer': 'none',
        'dropout_rate': 0.0,
        'lr_hparams': {
            'base_lr': 0.001,
            'schedule': 'constant'
        },
    })
    eval_num_batches = 5
    eval_batch_size = hps.batch_size
    loss_name = 'sigmoid_binary_cross_entropy'
    metrics_name = 'binary_classification_metrics'
    dataset, dataset_meta_data = _get_fake_graph_dataset(
        batch_size=hps.batch_size, eval_num_batches=eval_num_batches, hps=hps)
    model = model_cls(hps, dataset_meta_data, loss_name, metrics_name)
    initializer = initializers.get_initializer('noop')

    metrics_logger, init_logger = utils.set_up_loggers(self.test_dir)
    _ = list(

        trainer.train(