# 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/model_utils.py
# ACTIVATIONS = {
#     'relu': jax.nn.relu,
#     'id': lambda x: x,
#     'tanh': jnp.tanh,
#     'sigmoid': jax.nn.sigmoid,
#     'silu': jax.nn.silu,
#     'leaky_relu': jax.nn.leaky_relu,
# }
# INITIALIZERS = {
#     'delta_orthogonal': initializers.delta_orthogonal,
#     'orthogonal': initializers.orthogonal,
#     'lecun_normal': lecun_normal,
#     'xavier_uniform': initializers.xavier_uniform,
# }
# class ScalarMultiply(nn.Module):
#   def __call__(self, x):
# def get_normalizer(normalizer,
#                    train,
#                    batch_size=None,
#                    virtual_batch_size=None,
#                    total_batch_size=None,
#                    dtype=jnp.float32,
#                    data_format='NHWC'):
#     def identity_wrapper(*args, **kwargs):
#       def identity(x, *args, **kwargs):
# def apply_label_smoothing(one_hot_targets, label_smoothing):
# def sync_batchnorm_stats(state):
# def cross_device_avg(pytree):
# def l2_regularization(params, l2_decay_rank_threshold):
# def flatten_dict(nested_dict, sep='/'):
# def rescale_layers(params, layer_rescale_factors):


from typing import Any, Sequence
from flax import linen as nn
from init2winit.model_lib import base_model
from init2winit.model_lib import model_utils
from jax.nn import initializers
from ml_collections.config_dict import config_dict
import jax.numpy as jnp
  strides: Sequence[int]
  num_dense_units: int
  activation_fn: Any
  normalizer: str = 'none'
  kernel_init: model_utils.Initializer = initializers.lecun_normal()
  bias_init: model_utils.Initializer = initializers.zeros

  @nn.compact
  def __call__(self, x, train):
    maybe_normalize = model_utils.get_normalizer(self.normalizer, train)
    iterator = zip(
        self.num_filters, self.kernel_sizes, self.kernel_paddings,
        self.window_sizes, self.window_paddings, self.strides)
    for num_filters, kernel_size, kernel_padding, window_size, window_padding, stride in iterator:
      x = nn.Conv(
          num_filters, (kernel_size, kernel_size), (1, 1),
          padding=kernel_padding,
          kernel_init=self.kernel_init,
          bias_init=self.bias_init)(x)
      x = model_utils.ACTIVATIONS[self.activation_fn](x)
      x = maybe_normalize()(x)
      x = nn.max_pool(
          x,
          window_shape=(window_size, window_size),
          strides=(stride, stride),
          padding=window_padding)
    x = jnp.reshape(x, (x.shape[0], -1))
    for num_units in self.num_dense_units:
      x = nn.Dense(
          num_units, kernel_init=self.kernel_init, bias_init=self.bias_init)(x)

# 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/model_utils.py
# ACTIVATIONS = {
#     'relu': jax.nn.relu,
#     'id': lambda x: x,
#     'tanh': jnp.tanh,
#     'sigmoid': jax.nn.sigmoid,
#     'silu': jax.nn.silu,
#     'leaky_relu': jax.nn.leaky_relu,
# }
# INITIALIZERS = {
#     'delta_orthogonal': initializers.delta_orthogonal,
#     'orthogonal': initializers.orthogonal,
#     'lecun_normal': lecun_normal,
#     'xavier_uniform': initializers.xavier_uniform,
# }
# class ScalarMultiply(nn.Module):
#   def __call__(self, x):
# def get_normalizer(normalizer,
#                    train,
#                    batch_size=None,
#                    virtual_batch_size=None,
#                    total_batch_size=None,
#                    dtype=jnp.float32,
#                    data_format='NHWC'):
#     def identity_wrapper(*args, **kwargs):
#       def identity(x, *args, **kwargs):
# def apply_label_smoothing(one_hot_targets, label_smoothing):
# def sync_batchnorm_stats(state):
# def cross_device_avg(pytree):
# def l2_regularization(params, l2_decay_rank_threshold):
# def flatten_dict(nested_dict, sep='/'):
# def rescale_layers(params, layer_rescale_factors):


# Path: init2winit/model_lib/max_pooling_cnn.py
from typing import Any, Sequence
from flax import linen as nn
from init2winit.model_lib import base_model
from init2winit.model_lib import model_utils
from jax.nn import initializers
from ml_collections.config_dict import config_dict
import jax.numpy as jnp

  strides: Sequence[int]
  num_dense_units: int
  activation_fn: Any
  normalizer: str = 'none'
  kernel_init: model_utils.Initializer = initializers.lecun_normal()
  bias_init: model_utils.Initializer = initializers.zeros

  @nn.compact
  def __call__(self, x, train):
    maybe_normalize = model_utils.get_normalizer(self.normalizer, train)
    iterator = zip(
        self.num_filters, self.kernel_sizes, self.kernel_paddings,
        self.window_sizes, self.window_paddings, self.strides)
    for num_filters, kernel_size, kernel_padding, window_size, window_padding, stride in iterator:
      x = nn.Conv(
          num_filters, (kernel_size, kernel_size), (1, 1),
          padding=kernel_padding,
          kernel_init=self.kernel_init,
          bias_init=self.bias_init)(x)
      x = model_utils.ACTIVATIONS[self.activation_fn](x)
      x = maybe_normalize()(x)
      x = nn.max_pool(
          x,
          window_shape=(window_size, window_size),
          strides=(stride, stride),
          padding=window_padding)
    x = jnp.reshape(x, (x.shape[0], -1))
    for num_units in self.num_dense_units:
      x = nn.Dense(
          num_units, kernel_init=self.kernel_init, bias_init=self.bias_init)(x)

      x = model_utils.ACTIVATIONS[self.activation_fn](x)