# Path: init2winit/utils.py
# class TrainingDivergedError(Exception):
# class _MetricsLeafState:
# class TrainingMetricsGrabber:
# class MetricLogger(object):
# def dtype_from_str(dtype_string):
# def timed(f):
#   def wrapper(*args, **kwargs):
# def set_up_loggers(train_dir, xm_work_unit=None):
# def get_summary_tree(training_metrics_grabber):
#   def _reduce_node(node):
# def _update_param_stats(leaf_state, param_gradient, update, new_param, config):
# def _validate_config(config):
#   def create(model_params, config):
#     def _node_init(x):
#   def update(self, model_gradient, old_params, new_params):
#   def state_dict(self):
#   def restore_state(self, state, state_dict):
# def run_in_parallel(function, list_of_kwargs_to_function, num_workers):
# def should_eval(global_step, eval_frequency, eval_steps):
# def add_log_file(logfile):
# def load_pytrees(pytree_path):
#   def __init__(self,
#                csv_path='',
#                json_path='',
#                pytree_path='',
#                events_dir=None,
#                **logger_kwargs):
#   def append_scalar_metrics(self, metrics):
#   def write_pytree(self, pytree):
#   def append_pytree(self, pytree):
#   def append_json_object(self, json_obj):
# def _summary_str(param):
# def log_pytree_shape_and_statistics(pytree, json_path=None):
# 
# 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 jax
import numpy as np
from init2winit import utils
from init2winit.model_lib import model_utils
  # calls that would remove it), because we assume that there is no dependence
  # on the batch size with the model (batch norm reduces across a batch dim of
  # any size). This is similar to how the Flax examples initialize models:
  # https://github.com/google/flax/blob/44ee6f2f4130856d47159dc58981fb26ea2240f4/examples/imagenet/train.py#L70.
  if fake_batch:
    fake_input_batch = fake_batch
  elif isinstance(input_shape, list):  # Typical case for seq2seq models
    fake_input_batch = [
        np.zeros((2, *x), model_dtype) for x in input_shape
    ]
  else:  # Typical case for classification models
    fake_input_batch = [np.zeros((2, *input_shape), model_dtype)]
  params_rng, init_rng, dropout_rng = jax.random.split(rng, num=3)

  # By jitting the model init function, we initialize the model parameters
  # lazily without computing a full forward pass. For further documentation, see
  # https://flax.readthedocs.io/en/latest/flax.linen.html?highlight=jax.jit#flax.linen.Module.init.
  # We need to close over train=False here because otherwise the jitted init
  # function will convert the train Python bool to a jax boolean, which will
  # mess up Pythonic boolean statements like `not train` inside the model
  # construction.
  model_init_fn = jax.jit(functools.partial(flax_module.init, train=False))
  init_dict = model_init_fn(
      {'params': params_rng, 'dropout': dropout_rng},
      *fake_input_batch)
  # Trainable model parameters.
  params = init_dict['params']
  batch_stats = init_dict.get('batch_stats', {})

  if hps.get('layer_rescale_factors'):

# Path: init2winit/utils.py
# class TrainingDivergedError(Exception):
# class _MetricsLeafState:
# class TrainingMetricsGrabber:
# class MetricLogger(object):
# def dtype_from_str(dtype_string):
# def timed(f):
#   def wrapper(*args, **kwargs):
# def set_up_loggers(train_dir, xm_work_unit=None):
# def get_summary_tree(training_metrics_grabber):
#   def _reduce_node(node):
# def _update_param_stats(leaf_state, param_gradient, update, new_param, config):
# def _validate_config(config):
#   def create(model_params, config):
#     def _node_init(x):
#   def update(self, model_gradient, old_params, new_params):
#   def state_dict(self):
#   def restore_state(self, state, state_dict):
# def run_in_parallel(function, list_of_kwargs_to_function, num_workers):
# def should_eval(global_step, eval_frequency, eval_steps):
# def add_log_file(logfile):
# def load_pytrees(pytree_path):
#   def __init__(self,
#                csv_path='',
#                json_path='',
#                pytree_path='',
#                events_dir=None,
#                **logger_kwargs):
#   def append_scalar_metrics(self, metrics):
#   def write_pytree(self, pytree):
#   def append_pytree(self, pytree):
#   def append_json_object(self, json_obj):
# def _summary_str(param):
# def log_pytree_shape_and_statistics(pytree, json_path=None):
# 
# 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/init_lib/init_utils.py
import functools
import jax
import numpy as np
from init2winit import utils
from init2winit.model_lib import model_utils

  # calls that would remove it), because we assume that there is no dependence
  # on the batch size with the model (batch norm reduces across a batch dim of
  # any size). This is similar to how the Flax examples initialize models:
  # https://github.com/google/flax/blob/44ee6f2f4130856d47159dc58981fb26ea2240f4/examples/imagenet/train.py#L70.
  if fake_batch:
    fake_input_batch = fake_batch
  elif isinstance(input_shape, list):  # Typical case for seq2seq models
    fake_input_batch = [
        np.zeros((2, *x), model_dtype) for x in input_shape
    ]
  else:  # Typical case for classification models
    fake_input_batch = [np.zeros((2, *input_shape), model_dtype)]
  params_rng, init_rng, dropout_rng = jax.random.split(rng, num=3)

  # By jitting the model init function, we initialize the model parameters
  # lazily without computing a full forward pass. For further documentation, see
  # https://flax.readthedocs.io/en/latest/flax.linen.html?highlight=jax.jit#flax.linen.Module.init.
  # We need to close over train=False here because otherwise the jitted init
  # function will convert the train Python bool to a jax boolean, which will
  # mess up Pythonic boolean statements like `not train` inside the model
  # construction.
  model_init_fn = jax.jit(functools.partial(flax_module.init, train=False))
  init_dict = model_init_fn(
      {'params': params_rng, 'dropout': dropout_rng},
      *fake_input_batch)
  # Trainable model parameters.
  params = init_dict['params']
  batch_stats = init_dict.get('batch_stats', {})

  if hps.get('layer_rescale_factors'):

    params = model_utils.rescale_layers(params, hps.layer_rescale_factors)