# 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

    grad_loss = jax.grad(loss, has_aux=True)

    hessian = np.diag(
        np.array([
            1.0 / np.power(i, model_hps.hessian_decay_power)
            for i in range(1, dim + 1)
        ]))
    noise_matrix = np.diag(
        np.array([
            1.0 / np.power(i, model_hps.noise_decay_power / 2.0)
            for i in range(1, dim + 1)
        ]))

    noise = jnp.dot(noise_eps, noise_matrix)
    mean_noise = np.mean(noise, axis=0)

    # NQM gradient = Hx + eps   where eps ~ N(0, C / batch_size).
    expected_grad = np.dot(hessian, model_x) + mean_noise

    g = grad_loss(params, {'inputs': noise_eps})[0]['x']

    grad_error = np.sum(np.abs(g - expected_grad))
    self.assertAlmostEqual(grad_error, 0.0, places=5)

  @parameterized.named_parameters(*autoencoder_keys)
  def test_autoencoder_model(self, model_str):
    """Test forward pass of the autoencoder models."""

    model_cls = models.get_model(model_str)

# 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


    grad_loss = jax.grad(loss, has_aux=True)

    hessian = np.diag(
        np.array([
            1.0 / np.power(i, model_hps.hessian_decay_power)
            for i in range(1, dim + 1)
        ]))
    noise_matrix = np.diag(
        np.array([
            1.0 / np.power(i, model_hps.noise_decay_power / 2.0)
            for i in range(1, dim + 1)
        ]))

    noise = jnp.dot(noise_eps, noise_matrix)
    mean_noise = np.mean(noise, axis=0)

    # NQM gradient = Hx + eps   where eps ~ N(0, C / batch_size).
    expected_grad = np.dot(hessian, model_x) + mean_noise

    g = grad_loss(params, {'inputs': noise_eps})[0]['x']

    grad_error = np.sum(np.abs(g - expected_grad))
    self.assertAlmostEqual(grad_error, 0.0, places=5)

  @parameterized.named_parameters(*autoencoder_keys)
  def test_autoencoder_model(self, model_str):
    """Test forward pass of the autoencoder models."""

    model_cls = models.get_model(model_str)

    model_hps = models.get_model_hparams(model_str)