# 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):


import functools
import math
import jax.numpy as jnp
from flax import linen as nn
from init2winit.model_lib import base_model
from init2winit.model_lib import model_utils
from ml_collections.config_dict import config_dict
            'base_lr': 0.2,
        },
        normalizer='batch_norm',
        optimizer='momentum',
        opt_hparams={
            'momentum': 0.9,
        },
        batch_size=128,
        l2_decay_factor=0.0001,
        l2_decay_rank_threshold=2,
        label_smoothing=None,
        rng_seed=-1,
        use_shallue_label_smoothing=False,
        model_dtype='float32',
        grad_clip=None,
        normalize_classifier_input='none',
        classification_scale_factor=1.0,
    ))


class BottleneckBlock(nn.Module):
  """Composition function with a bottleneck layer.

  The composition function, which is the basic unit of a dense block, conists
  of a batch normalization, followed by a ReLU and 3x3 convolution.  In order
  to reduce the number of input feature maps, the above operations are
  preceded by 1x1 convoluational operation (and the correpsonding batch
  normalization and ReLU).
  """
  growth_rate: int

# 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/adabelief_densenet.py
import functools
import math
import jax.numpy as jnp
from flax import linen as nn
from init2winit.model_lib import base_model
from init2winit.model_lib import model_utils
from ml_collections.config_dict import config_dict

            'base_lr': 0.2,
        },
        normalizer='batch_norm',
        optimizer='momentum',
        opt_hparams={
            'momentum': 0.9,
        },
        batch_size=128,
        l2_decay_factor=0.0001,
        l2_decay_rank_threshold=2,
        label_smoothing=None,
        rng_seed=-1,
        use_shallue_label_smoothing=False,
        model_dtype='float32',
        grad_clip=None,
        normalize_classifier_input='none',
        classification_scale_factor=1.0,
    ))


class BottleneckBlock(nn.Module):
  """Composition function with a bottleneck layer.

  The composition function, which is the basic unit of a dense block, conists
  of a batch normalization, followed by a ReLU and 3x3 convolution.  In order
  to reduce the number of input feature maps, the above operations are
  preceded by 1x1 convoluational operation (and the correpsonding batch
  normalization and ReLU).
  """
  growth_rate: int

  dtype: model_utils.Dtype = jnp.float32