# Path: init2winit/init_lib/meta_init.py
# def meta_init(loss_fn,
#               flax_module,
#               params,
#               hps,
#               input_shape,
#               output_shape,
#               rng_key,
#               metrics_logger=None,
#               log_every=10):
#   """Implements MetaInit initializer.
# 
#   Args:
#     loss_fn: Loss function.
#     flax_module: Flax nn.Module class.
#     params: The dict of model parameters.
#     hps: HParam object. Required hparams are meta_learning_rate,
#       meta_batch_size, meta_steps, and epsilon.
#     input_shape: Must agree with batch[0].shape[1:].
#     output_shape: Must agree with batch[1].shape[1:].
#     rng_key: jax.PRNGKey, used to seed all randomness.
#     metrics_logger: Instance of utils.MetricsLogger
#     log_every: Print meta loss every k steps.
# 
#   Returns:
#     A Flax module with the learned initialization.
#   """
#   # Pretty print the preinitialized norms with the variable shapes.
#   if jax.process_index() == 0:
#     logging.info('Preinitialized norms:')
#     _log_shape_and_norms(params, metrics_logger, key='init_norms')
#     # First grab the norms of all weights and rescale params to have norm 1.
#     logging.info('Running meta init')
#   norms = jax.tree_map(lambda node: jnp.linalg.norm(node.reshape(-1)),
#                        params)
# 
#   normalized_params = jax.tree_map(normalize, params)
# 
#   learned_norms, _ = meta_optimize_scales(
#       loss_fn,
#       flax_module.apply,
#       normalized_params,
#       norms,
#       hps,
#       input_shape,
#       output_shape,
#       rng_key,
#       metrics_logger=metrics_logger,
#       log_every=log_every)
#   new_params = scale_params(normalized_params, learned_norms)
# 
#   if jax.process_index() == 0:
#     # Pretty print the meta init norms with the variable shapes.
#     logging.info('Learned norms from meta_init:')
#     _log_shape_and_norms(new_params, metrics_logger, key='meta_init_norms')
# 
#   return new_params
# 
# Path: init2winit/init_lib/sparse_init.py
# def sparse_init(loss_fn,
#                 flax_module,
#                 params,
#                 hps,
#                 input_shape,
#                 output_shape,
#                 rng_key,
#                 metrics_logger=None,
#                 log_every=10):
#   """Implements SparseInit initializer.
# 
#   Args:
#     loss_fn: Loss function.
#     flax_module: Flax nn.Module class.
#     params: The dict of model parameters.
#     hps: HParam object. Required hparams are meta_learning_rate,
#       meta_batch_size, meta_steps, and epsilon.
#     input_shape: Must agree with batch[0].shape[1:].
#     output_shape: Must agree with batch[1].shape[1:].
#     rng_key: jax.PRNGKey, used to seed all randomness.
#     metrics_logger: Instance of utils.MetricsLogger
#     log_every: Print meta loss every k steps.
# 
#   Returns:
#     A Flax model with sparse initialization.
#   """
# 
#   del flax_module, loss_fn, input_shape, output_shape, rng_key, metrics_logger, log_every
# 
#   params = unfreeze(params)
#   activation_functions = hps.activation_function
#   num_hidden_layers = len(hps.hid_sizes)
#   if isinstance(hps.activation_function, str):
#     activation_functions = [hps.activation_function] * num_hidden_layers
#   for i, key in enumerate(params):
#     num_units, num_weights = params[key]['kernel'].shape
#     mask = np.zeros((num_units, num_weights), dtype=bool)
#     for k in range(num_units):
#       if num_weights >= hps.non_zero_connection_weights:
#         sample = np.random.choice(
#             num_weights, hps.non_zero_connection_weights, replace=False)
#       else:
#         sample = np.random.choice(num_weights, hps.non_zero_connection_weights)
#       mask[k, sample] = True
#     params[key]['kernel'] = params[key]['kernel'].at[~mask].set(0.0)
#     if i < num_hidden_layers and activation_functions[i] == 'tanh':
#       params[key]['bias'] = params[key]['bias'].at[:].set(0.5)
#     else:
#       params[key]['bias'] = params[key]['bias'].at[:].set(0.0)
#   return frozen_dict.freeze(params)


from init2winit.init_lib import meta_init
from init2winit.init_lib import sparse_init
from ml_collections.config_dict import config_dict
new_params = init(loss, init_params, hps, num_outputs, input_shape)

TODO(gilmer, gdahl, schsam): The API of an initializer should in general be
aware of the moments of the data. Currently we assume that all input coordinates
are iid standard normal distributions.
"""



# This function is meant to match the general API of an initializer
# pylint: disable=unused-argument
def noop(
    loss_fn=None,
    flax_module=None,
    params=None,
    hps=None,
    input_shape=None,
    output_shape=None,
    rng_key=None,
    metrics_logger=None,
):
  """No-op init."""
  return params
# pylint: enable=unused-argument

DEFAULT_HPARAMS = config_dict.ConfigDict()

_ALL_INITIALIZERS = {
    'noop': (noop, DEFAULT_HPARAMS),
    'meta_init': (meta_init.meta_init, meta_init.DEFAULT_HPARAMS),

# Path: init2winit/init_lib/meta_init.py
# def meta_init(loss_fn,
#               flax_module,
#               params,
#               hps,
#               input_shape,
#               output_shape,
#               rng_key,
#               metrics_logger=None,
#               log_every=10):
#   """Implements MetaInit initializer.
# 
#   Args:
#     loss_fn: Loss function.
#     flax_module: Flax nn.Module class.
#     params: The dict of model parameters.
#     hps: HParam object. Required hparams are meta_learning_rate,
#       meta_batch_size, meta_steps, and epsilon.
#     input_shape: Must agree with batch[0].shape[1:].
#     output_shape: Must agree with batch[1].shape[1:].
#     rng_key: jax.PRNGKey, used to seed all randomness.
#     metrics_logger: Instance of utils.MetricsLogger
#     log_every: Print meta loss every k steps.
# 
#   Returns:
#     A Flax module with the learned initialization.
#   """
#   # Pretty print the preinitialized norms with the variable shapes.
#   if jax.process_index() == 0:
#     logging.info('Preinitialized norms:')
#     _log_shape_and_norms(params, metrics_logger, key='init_norms')
#     # First grab the norms of all weights and rescale params to have norm 1.
#     logging.info('Running meta init')
#   norms = jax.tree_map(lambda node: jnp.linalg.norm(node.reshape(-1)),
#                        params)
# 
#   normalized_params = jax.tree_map(normalize, params)
# 
#   learned_norms, _ = meta_optimize_scales(
#       loss_fn,
#       flax_module.apply,
#       normalized_params,
#       norms,
#       hps,
#       input_shape,
#       output_shape,
#       rng_key,
#       metrics_logger=metrics_logger,
#       log_every=log_every)
#   new_params = scale_params(normalized_params, learned_norms)
# 
#   if jax.process_index() == 0:
#     # Pretty print the meta init norms with the variable shapes.
#     logging.info('Learned norms from meta_init:')
#     _log_shape_and_norms(new_params, metrics_logger, key='meta_init_norms')
# 
#   return new_params
# 
# Path: init2winit/init_lib/sparse_init.py
# def sparse_init(loss_fn,
#                 flax_module,
#                 params,
#                 hps,
#                 input_shape,
#                 output_shape,
#                 rng_key,
#                 metrics_logger=None,
#                 log_every=10):
#   """Implements SparseInit initializer.
# 
#   Args:
#     loss_fn: Loss function.
#     flax_module: Flax nn.Module class.
#     params: The dict of model parameters.
#     hps: HParam object. Required hparams are meta_learning_rate,
#       meta_batch_size, meta_steps, and epsilon.
#     input_shape: Must agree with batch[0].shape[1:].
#     output_shape: Must agree with batch[1].shape[1:].
#     rng_key: jax.PRNGKey, used to seed all randomness.
#     metrics_logger: Instance of utils.MetricsLogger
#     log_every: Print meta loss every k steps.
# 
#   Returns:
#     A Flax model with sparse initialization.
#   """
# 
#   del flax_module, loss_fn, input_shape, output_shape, rng_key, metrics_logger, log_every
# 
#   params = unfreeze(params)
#   activation_functions = hps.activation_function
#   num_hidden_layers = len(hps.hid_sizes)
#   if isinstance(hps.activation_function, str):
#     activation_functions = [hps.activation_function] * num_hidden_layers
#   for i, key in enumerate(params):
#     num_units, num_weights = params[key]['kernel'].shape
#     mask = np.zeros((num_units, num_weights), dtype=bool)
#     for k in range(num_units):
#       if num_weights >= hps.non_zero_connection_weights:
#         sample = np.random.choice(
#             num_weights, hps.non_zero_connection_weights, replace=False)
#       else:
#         sample = np.random.choice(num_weights, hps.non_zero_connection_weights)
#       mask[k, sample] = True
#     params[key]['kernel'] = params[key]['kernel'].at[~mask].set(0.0)
#     if i < num_hidden_layers and activation_functions[i] == 'tanh':
#       params[key]['bias'] = params[key]['bias'].at[:].set(0.5)
#     else:
#       params[key]['bias'] = params[key]['bias'].at[:].set(0.0)
#   return frozen_dict.freeze(params)


# Path: init2winit/init_lib/initializers.py
from init2winit.init_lib import meta_init
from init2winit.init_lib import sparse_init
from ml_collections.config_dict import config_dict

new_params = init(loss, init_params, hps, num_outputs, input_shape)

TODO(gilmer, gdahl, schsam): The API of an initializer should in general be
aware of the moments of the data. Currently we assume that all input coordinates
are iid standard normal distributions.
"""



# This function is meant to match the general API of an initializer
# pylint: disable=unused-argument
def noop(
    loss_fn=None,
    flax_module=None,
    params=None,
    hps=None,
    input_shape=None,
    output_shape=None,
    rng_key=None,
    metrics_logger=None,
):
  """No-op init."""
  return params
# pylint: enable=unused-argument

DEFAULT_HPARAMS = config_dict.ConfigDict()

_ALL_INITIALIZERS = {
    'noop': (noop, DEFAULT_HPARAMS),
    'meta_init': (meta_init.meta_init, meta_init.DEFAULT_HPARAMS),

    'sparse_init': (sparse_init.sparse_init, sparse_init.DEFAULT_HPARAMS),