# 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
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for checkpoint.py."""



FLAGS = flags.FLAGS


def pytree_equal(tree1, tree2):
  try:
    equal_tree = jax.tree_util.tree_multimap(np.array_equal, tree1, tree2)
    return jax.tree_util.tree_reduce(lambda x, y: x and y, equal_tree)
  # The tree_utils will raise TypeErrors if structures don't match.
  except TypeError:
    return False

INPUT_SHAPE = [10, 28, 28, 1]
OUTPUT_SHAPE = (10,)


class CheckpointTest(parameterized.TestCase):
  """Tests for checkpoint.py mostly by doing round trips."""

  def setUp(self):
    super(CheckpointTest, self).setUp()
    self.test_dir = tempfile.mkdtemp()
    loss_name = 'cross_entropy'
    metrics_name = 'classification_metrics'

# 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

# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for checkpoint.py."""



FLAGS = flags.FLAGS


def pytree_equal(tree1, tree2):
  try:
    equal_tree = jax.tree_util.tree_multimap(np.array_equal, tree1, tree2)
    return jax.tree_util.tree_reduce(lambda x, y: x and y, equal_tree)
  # The tree_utils will raise TypeErrors if structures don't match.
  except TypeError:
    return False

INPUT_SHAPE = [10, 28, 28, 1]
OUTPUT_SHAPE = (10,)


class CheckpointTest(parameterized.TestCase):
  """Tests for checkpoint.py mostly by doing round trips."""

  def setUp(self):
    super(CheckpointTest, self).setUp()
    self.test_dir = tempfile.mkdtemp()
    loss_name = 'cross_entropy'
    metrics_name = 'classification_metrics'

    model = models.get_model('fully_connected')