# Path: init2winit/checkpoint.py
# 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_'):
# 
# Path: init2winit/model_lib/models.py
# _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):


import copy
import functools
import os.path
import shutil
import tempfile
import jax.numpy as jnp
import jax.tree_util
import numpy as np
from absl import flags
from absl.testing import absltest
from absl.testing import parameterized
from init2winit import checkpoint
from init2winit.model_lib import models
from tensorflow.io import gfile
  def test_delete_old_checkpoints(self):
    """Test that old checkpoints are deleted."""
    state1 = dict(params=self.params,
                  global_step=5,
                  completed_epochs=4)
    checkpoint.save_checkpoint(
        self.test_dir,
        0,
        state1,
        max_to_keep=1)

    state2 = dict(params=self.params,
                  global_step=10,
                  completed_epochs=8)
    checkpoint.save_checkpoint(
        self.test_dir,
        1,
        state2,
        max_to_keep=1)
    dir_contents = gfile.glob(os.path.join(self.test_dir, '*'))
    self.assertLen(dir_contents, 1)

  def test_save_checkpoint_background_reraises_error(self):
    """Test than an error while saving a checkpoint is re-raised later."""
    # Checkpoint error is not raised when it actually happens, but when we next
    # write a checkpoint.
    baz = ['a', 'b', 'ccc']
    state = dict(params=self.params,
                 global_step=5, completed_epochs=4,
                 baz=baz)

# Path: init2winit/checkpoint.py
# 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_'):
# 
# Path: init2winit/model_lib/models.py
# _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):


# Path: init2winit/test_checkpoint.py
import copy
import functools
import os.path
import shutil
import tempfile
import jax.numpy as jnp
import jax.tree_util
import numpy as np
from absl import flags
from absl.testing import absltest
from absl.testing import parameterized
from init2winit import checkpoint
from init2winit.model_lib import models
from tensorflow.io import gfile

  def test_delete_old_checkpoints(self):
    """Test that old checkpoints are deleted."""
    state1 = dict(params=self.params,
                  global_step=5,
                  completed_epochs=4)
    checkpoint.save_checkpoint(
        self.test_dir,
        0,
        state1,
        max_to_keep=1)

    state2 = dict(params=self.params,
                  global_step=10,
                  completed_epochs=8)
    checkpoint.save_checkpoint(
        self.test_dir,
        1,
        state2,
        max_to_keep=1)
    dir_contents = gfile.glob(os.path.join(self.test_dir, '*'))
    self.assertLen(dir_contents, 1)

  def test_save_checkpoint_background_reraises_error(self):
    """Test than an error while saving a checkpoint is re-raised later."""
    # Checkpoint error is not raised when it actually happens, but when we next
    # write a checkpoint.
    baz = ['a', 'b', 'ccc']
    state = dict(params=self.params,
                 global_step=5, completed_epochs=4,
                 baz=baz)

    checkpoint.save_checkpoint_background(