# Path: init2winit/model_lib/base_model.py
# 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):
# 
# Path: init2winit/model_lib/fully_connected.py
# class FullyConnected(nn.Module):
#   """Defines a fully connected neural network.
# 
#   The model assumes the input data has shape
#   [batch_size_per_device, *input_shape] where input_shape may be of arbitrary
#   rank. The model flatten the input before applying a dense layer.
#   """
#   num_outputs: int
#   hid_sizes: Tuple[int]
#   activation_function: Any
#   kernel_inits: Tuple[model_utils.Initializer]
#   bias_init: model_utils.Initializer = initializers.zeros
# 
#   @nn.compact
#   def __call__(self, x, train):
#     del train
#     if not isinstance(self.activation_function, str):
#       if len(self.activation_function) != len(self.hid_sizes):
#         raise ValueError(
#             'The number of activation functions must be equal to the number '
#             'of hidden layers')
#       activation_function = copy.deepcopy(self.activation_function)
#     else:
#       activation_function = [self.activation_function] * len(self.hid_sizes)
# 
#     x = jnp.reshape(x, (x.shape[0], -1))
#     for i, (num_hid, init) in enumerate(
#         zip(self.hid_sizes, self.kernel_inits[:-1])):
#       x = nn.Dense(num_hid, kernel_init=init, bias_init=self.bias_init)(x)
#       x = model_utils.ACTIVATIONS[activation_function[i]](x)
#     x = nn.Dense(
#         self.num_outputs,
#         kernel_init=self.kernel_inits[-1],
#         bias_init=self.bias_init)(x)
#     return x


from init2winit.model_lib import base_model
from init2winit.model_lib.fully_connected import FullyConnected
from jax.nn import initializers
from ml_collections.config_dict import config_dict
        hid_sizes=[128, 64, 32, 64, 128],
        activation_function=['relu', 'relu', 'relu', 'relu', 'relu'],
        kernel_scales=[0.5, 0.5, 0.5, 0.5, 0.5, 0.5],
        lr_hparams={
            'base_lr': 0.1,
            'schedule': 'constant'
        },
        layer_rescale_factors={},
        optimizer='hessian_free',
        opt_hparams={
            'cg_max_iters': 250,
        },
        batch_size=128,
        l2_decay_factor=None,
        l2_decay_rank_threshold=0,
        label_smoothing=None,
        rng_seed=-1,
        use_shallue_label_smoothing=False,
        model_dtype='float32',
    ))


class AutoEncoderModel(base_model.BaseModel):

  def build_flax_module(self):
    kernel_inits = [
        initializers.variance_scaling(scale, 'fan_in', 'truncated_normal')
        for scale in self.hps.kernel_scales
    ]


# Path: init2winit/model_lib/base_model.py
# 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):
# 
# Path: init2winit/model_lib/fully_connected.py
# class FullyConnected(nn.Module):
#   """Defines a fully connected neural network.
# 
#   The model assumes the input data has shape
#   [batch_size_per_device, *input_shape] where input_shape may be of arbitrary
#   rank. The model flatten the input before applying a dense layer.
#   """
#   num_outputs: int
#   hid_sizes: Tuple[int]
#   activation_function: Any
#   kernel_inits: Tuple[model_utils.Initializer]
#   bias_init: model_utils.Initializer = initializers.zeros
# 
#   @nn.compact
#   def __call__(self, x, train):
#     del train
#     if not isinstance(self.activation_function, str):
#       if len(self.activation_function) != len(self.hid_sizes):
#         raise ValueError(
#             'The number of activation functions must be equal to the number '
#             'of hidden layers')
#       activation_function = copy.deepcopy(self.activation_function)
#     else:
#       activation_function = [self.activation_function] * len(self.hid_sizes)
# 
#     x = jnp.reshape(x, (x.shape[0], -1))
#     for i, (num_hid, init) in enumerate(
#         zip(self.hid_sizes, self.kernel_inits[:-1])):
#       x = nn.Dense(num_hid, kernel_init=init, bias_init=self.bias_init)(x)
#       x = model_utils.ACTIVATIONS[activation_function[i]](x)
#     x = nn.Dense(
#         self.num_outputs,
#         kernel_init=self.kernel_inits[-1],
#         bias_init=self.bias_init)(x)
#     return x


# Path: init2winit/model_lib/autoencoder.py
from init2winit.model_lib import base_model
from init2winit.model_lib.fully_connected import FullyConnected
from jax.nn import initializers
from ml_collections.config_dict import config_dict

        hid_sizes=[128, 64, 32, 64, 128],
        activation_function=['relu', 'relu', 'relu', 'relu', 'relu'],
        kernel_scales=[0.5, 0.5, 0.5, 0.5, 0.5, 0.5],
        lr_hparams={
            'base_lr': 0.1,
            'schedule': 'constant'
        },
        layer_rescale_factors={},
        optimizer='hessian_free',
        opt_hparams={
            'cg_max_iters': 250,
        },
        batch_size=128,
        l2_decay_factor=None,
        l2_decay_rank_threshold=0,
        label_smoothing=None,
        rng_seed=-1,
        use_shallue_label_smoothing=False,
        model_dtype='float32',
    ))


class AutoEncoderModel(base_model.BaseModel):

  def build_flax_module(self):
    kernel_inits = [
        initializers.variance_scaling(scale, 'fan_in', 'truncated_normal')
        for scale in self.hps.kernel_scales
    ]


    return FullyConnected(