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

  # Again average across all of the shards.
  hgp = model_utils.cross_device_avg(jvp(grad_loss, [scaled_params], [g])[1])

  nparams = _count_params(g)

  def meta_term(g, hgp):
    ratio = (g-hgp) / (g + epsilon * jax.lax.stop_gradient(2*(g >= 0) - 1))
    return jnp.sum(jnp.abs(ratio - 1))

  return jax.tree_util.tree_reduce(
      operator.add, jax.tree_multimap(meta_term, g, hgp)) / nparams


def normalize(node):
  norm = jnp.linalg.norm(node.reshape(-1))
  if jnp.abs(norm) == 0.0:
    return node
  return node / norm


def _get_non_bias_params(params):

# 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

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

  # Again average across all of the shards.
  hgp = model_utils.cross_device_avg(jvp(grad_loss, [scaled_params], [g])[1])

  nparams = _count_params(g)

  def meta_term(g, hgp):
    ratio = (g-hgp) / (g + epsilon * jax.lax.stop_gradient(2*(g >= 0) - 1))
    return jnp.sum(jnp.abs(ratio - 1))

  return jax.tree_util.tree_reduce(
      operator.add, jax.tree_multimap(meta_term, g, hgp)) / nparams


def normalize(node):
  norm = jnp.linalg.norm(node.reshape(-1))
  if jnp.abs(norm) == 0.0:
    return node
  return node / norm


def _get_non_bias_params(params):

  flat_params = model_utils.flatten_dict(params)