# 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 jax.numpy as jnp
import jax.tree_util
import jraph
import numpy as np
from absl.testing import absltest
from absl.testing import parameterized
from init2winit.model_lib import models
from jax.flatten_util import ravel_pytree
from ml_collections.config_dict import config_dict

"""




HIDDEN_SIZES = (50, 50)

INPUT_SHAPE = {
    'classification': (64, 32, 32, 3),
    'autoencoder': (64, 28, 28, 1),
    'convolutional_autoencoder': (64, 28, 28, 1),
}
OUTPUT_SHAPE = {
    'classification': (5,),
    'autoencoder': (784,),
    'convolutional_autoencoder': (28, 28, 1),
}

# Automatically test all defined models.
autoencoder_models = ['autoencoder', 'convolutional_autoencoder']
text_models = ['transformer', 'performer']
seq2seq_models = ['xformer_translate']

autoencoder_keys = [('test_{}'.format(m), m) for m in autoencoder_models]
excluded_classification_models = text_models + autoencoder_models + seq2seq_models + [
    'nqm', 'gnn', 'dlrm', 'vit'
]
classification_keys = [
    ('test_{}'.format(m), m)

# 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/model_lib/test_models.py
import copy
import functools
import jax.numpy as jnp
import jax.tree_util
import jraph
import numpy as np
from absl.testing import absltest
from absl.testing import parameterized
from init2winit.model_lib import models
from jax.flatten_util import ravel_pytree
from ml_collections.config_dict import config_dict


"""




HIDDEN_SIZES = (50, 50)

INPUT_SHAPE = {
    'classification': (64, 32, 32, 3),
    'autoencoder': (64, 28, 28, 1),
    'convolutional_autoencoder': (64, 28, 28, 1),
}
OUTPUT_SHAPE = {
    'classification': (5,),
    'autoencoder': (784,),
    'convolutional_autoencoder': (28, 28, 1),
}

# Automatically test all defined models.
autoencoder_models = ['autoencoder', 'convolutional_autoencoder']
text_models = ['transformer', 'performer']
seq2seq_models = ['xformer_translate']

autoencoder_keys = [('test_{}'.format(m), m) for m in autoencoder_models]
excluded_classification_models = text_models + autoencoder_models + seq2seq_models + [
    'nqm', 'gnn', 'dlrm', 'vit'
]
classification_keys = [
    ('test_{}'.format(m), m)

    for m in models._ALL_MODELS.keys()  # pylint: disable=protected-access