# 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/dataset_lib/data_utils.py
# def iterator_as_numpy(iterator):
# def image_iterator(data,
#                    rescale,
#                    output_shape,
#                    is_one_hot,
#                    autoencoder,
#                    shuffle_rng=None,
#                    augment_fn=None,
#                    include_example_keys=False):
# def maybe_pad_batch(batch,
#                     desired_batch_size,
#                     data_format=None,
#                     mask_key=None):
#   def zero_pad(ar, pad_axis):
# def shard(batch, n_devices=None):
#   def _shard_array(array):
# def tf_to_numpy(tfds_data):


import functools
import itertools
import os
import jax
import jax.numpy as jnp
import numpy as np
from init2winit import utils
from init2winit.dataset_lib import data_utils
from init2winit.hessian import model_debugger
  rng = jax.random.fold_in(rng, local_device_index)

  def opt_cost(flax_module):
    return training_cost(flax_module, batch, batch_stats, rng)

  grad_fn = jax.value_and_grad(opt_cost, has_aux=True)
  (cost_value, _), grad = grad_fn(params)

  cost_value, grad = jax.lax.pmean((cost_value, grad), axis_name='batch')
  grad_norms = jax.tree_map(lambda x: jnp.linalg.norm(x.reshape(-1))**2, grad)
  param_norms = jax.tree_map(lambda x: jnp.linalg.norm(x.reshape(-1))**2,
                             params)
  return param_norms, grad_norms


class ModelDebugCallback:
  """Used to run the hessian eval in the trainer binary."""

  def __init__(self, model, optimizer, batch_stats, dataset, hps,
               callback_config, train_dir, rng):
    del hps
    del rng
    del optimizer
    del batch_stats
    del callback_config  # In future CL's we will use this.
    checkpoint_dir = os.path.join(train_dir, 'checkpoints')
    # copy batch_stats as we close over it, and it gets modified.
    self.dataset = dataset
    checkpoint_dir = os.path.join(train_dir, 'checkpoints')
    pytree_path = os.path.join(checkpoint_dir, 'debugger')

# 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/dataset_lib/data_utils.py
# def iterator_as_numpy(iterator):
# def image_iterator(data,
#                    rescale,
#                    output_shape,
#                    is_one_hot,
#                    autoencoder,
#                    shuffle_rng=None,
#                    augment_fn=None,
#                    include_example_keys=False):
# def maybe_pad_batch(batch,
#                     desired_batch_size,
#                     data_format=None,
#                     mask_key=None):
#   def zero_pad(ar, pad_axis):
# def shard(batch, n_devices=None):
#   def _shard_array(array):
# def tf_to_numpy(tfds_data):


# Path: hessian/model_debugger_callback.py
import functools
import itertools
import os
import jax
import jax.numpy as jnp
import numpy as np
from init2winit import utils
from init2winit.dataset_lib import data_utils
from init2winit.hessian import model_debugger

  rng = jax.random.fold_in(rng, local_device_index)

  def opt_cost(flax_module):
    return training_cost(flax_module, batch, batch_stats, rng)

  grad_fn = jax.value_and_grad(opt_cost, has_aux=True)
  (cost_value, _), grad = grad_fn(params)

  cost_value, grad = jax.lax.pmean((cost_value, grad), axis_name='batch')
  grad_norms = jax.tree_map(lambda x: jnp.linalg.norm(x.reshape(-1))**2, grad)
  param_norms = jax.tree_map(lambda x: jnp.linalg.norm(x.reshape(-1))**2,
                             params)
  return param_norms, grad_norms


class ModelDebugCallback:
  """Used to run the hessian eval in the trainer binary."""

  def __init__(self, model, optimizer, batch_stats, dataset, hps,
               callback_config, train_dir, rng):
    del hps
    del rng
    del optimizer
    del batch_stats
    del callback_config  # In future CL's we will use this.
    checkpoint_dir = os.path.join(train_dir, 'checkpoints')
    # copy batch_stats as we close over it, and it gets modified.
    self.dataset = dataset
    checkpoint_dir = os.path.join(train_dir, 'checkpoints')
    pytree_path = os.path.join(checkpoint_dir, 'debugger')

    logger = utils.MetricLogger(pytree_path=pytree_path)