# 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 json
import operator
import time
import jax
import jax.numpy as jnp
import numpy as np
import optax
from absl import logging
from flax import jax_utils
from flax import optim as optimizers
from flax.core import unfreeze
from init2winit.model_lib import model_utils
from jax import jvp
from ml_collections.config_dict import config_dict


def meta_loss(params_to_loss, scalars, normalized_params, epsilon):
  """Computes the meta_init objective function.

  This is a function of the learned scalars (one for each non-bias term in the
  model), and the original loss function. Note, the meta_loss is only a function
  of the variables with ndim > 1. Be careful how variables are shaped, this is
  potentially brittle, but the easiest way to avoid optimizing the bias terms.

  Args:
    params_to_loss: Function mapping (params) -> scalar valued loss.
    scalars: A pytree of scalars, one for each variable in the model. Note, we
      assume that all bias terms have a scale of 0.
    normalized_params: The pytree of model params. Requires that all params have
      norm equal to 1.
    epsilon: Used for numerical stability to avoid dividing by 0.

  Returns:
    Scalar value of the meta objective.
  """

  # Note, params_to_loss has captured the data for this particular shard.
  grad_loss = jax.grad(params_to_loss)

  scaled_params = scale_params(normalized_params, scalars)

  # The meta loss is a function of the data across all shards and does NOT
  # decompose as a sum over each data point. However the inner loss DOES
  # decompose, so we need to average the gradients of all of the data here.

# 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/meta_init.py
import functools
import json
import operator
import time
import jax
import jax.numpy as jnp
import numpy as np
import optax
from absl import logging
from flax import jax_utils
from flax import optim as optimizers
from flax.core import unfreeze
from init2winit.model_lib import model_utils
from jax import jvp
from ml_collections.config_dict import config_dict



def meta_loss(params_to_loss, scalars, normalized_params, epsilon):
  """Computes the meta_init objective function.

  This is a function of the learned scalars (one for each non-bias term in the
  model), and the original loss function. Note, the meta_loss is only a function
  of the variables with ndim > 1. Be careful how variables are shaped, this is
  potentially brittle, but the easiest way to avoid optimizing the bias terms.

  Args:
    params_to_loss: Function mapping (params) -> scalar valued loss.
    scalars: A pytree of scalars, one for each variable in the model. Note, we
      assume that all bias terms have a scale of 0.
    normalized_params: The pytree of model params. Requires that all params have
      norm equal to 1.
    epsilon: Used for numerical stability to avoid dividing by 0.

  Returns:
    Scalar value of the meta objective.
  """

  # Note, params_to_loss has captured the data for this particular shard.
  grad_loss = jax.grad(params_to_loss)

  scaled_params = scale_params(normalized_params, scalars)

  # The meta loss is a function of the data across all shards and does NOT
  # decompose as a sum over each data point. However the inner loss DOES
  # decompose, so we need to average the gradients of all of the data here.

  g = model_utils.cross_device_avg(grad_loss(scaled_params))