# Path: init2winit/checkpoint.py
# FLAGS = flags.FLAGS
# def replicate_checkpoint(
#     latest,
#     pytree_keys: Sequence[str],
#     replicate=True):
# def replicate_and_maybe_restore_latest_checkpoint(
#     unreplicated_optimizer_state,
#     unreplicated_params,
#     unreplicated_batch_stats,
#     unreplicated_training_metrics_grabber,
#     train_dir):
# def save_unreplicated_checkpoint_background(
#     train_dir,
#     optimizer_state,
#     params,
#     batch_stats,
#     training_metrics_grabber,
#     global_step,
#     preemption_count,
#     sum_train_cost,
#     max_to_keep=1):
# def _save_checkpoint_background_catch_error(*args, **kwargs):
# def wait_for_checkpoint_save():
# def save_checkpoint_background(*args, **kwargs):
# def save_checkpoint(train_dir,
#                     step,
#                     state,
#                     prefix='ckpt_',
#                     max_to_keep=None):
# def load_checkpoint(
#     checkpoint_path,
#     target=None,
#     prefix='ckpt_'):
# def load_latest_checkpoint(train_dir, target=None, prefix='ckpt_'):
# 
# 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 concurrent.futures
import functools
import json
import logging
import operator
import os.path
import time
import jax
import jax.numpy as jnp
import numpy as np
import pandas as pd
from typing import Any, Dict
from absl import logging as absl_logging
from clu import metric_writers
from flax import serialization
from flax import struct
from flax.training import checkpoints as flax_checkpoints
from init2winit import checkpoint
from init2winit.model_lib import model_utils
from tensorflow.io import gfile

  init_csv_path = os.path.join(train_dir, 'init_measurements.csv')
  init_json_path = os.path.join(train_dir, 'init_scalars.json')
  init_logger = MetricLogger(
      csv_path=init_csv_path,
      json_path=init_json_path,
      xm_work_unit=xm_work_unit)
  return metrics_logger, init_logger


def get_summary_tree(training_metrics_grabber):
  """Extracts desired training statistics from the grabber state.

  Currently this function will compute the scalar aggregate gradient variance
  for every weight matrix of the model. Future iterations of this function
  may depend on the metrics_grabber config.

  Args:
    training_metrics_grabber: TrainingMetricsGrabber object.

  Returns:
    A dict of different aggregate training statistics.
  """
  unreplicated_metrics_tree = jax.tree_map(lambda x: x[0],
                                           training_metrics_grabber.state)

  # Example key: Layer1/conv1/kernel/
  # NOTE: jax.tree_map does not work here, because tree_map will additionally
  # flatten the node state, while model_utils.flatten_dict will consider the
  # node object a leaf.

# Path: init2winit/checkpoint.py
# FLAGS = flags.FLAGS
# def replicate_checkpoint(
#     latest,
#     pytree_keys: Sequence[str],
#     replicate=True):
# def replicate_and_maybe_restore_latest_checkpoint(
#     unreplicated_optimizer_state,
#     unreplicated_params,
#     unreplicated_batch_stats,
#     unreplicated_training_metrics_grabber,
#     train_dir):
# def save_unreplicated_checkpoint_background(
#     train_dir,
#     optimizer_state,
#     params,
#     batch_stats,
#     training_metrics_grabber,
#     global_step,
#     preemption_count,
#     sum_train_cost,
#     max_to_keep=1):
# def _save_checkpoint_background_catch_error(*args, **kwargs):
# def wait_for_checkpoint_save():
# def save_checkpoint_background(*args, **kwargs):
# def save_checkpoint(train_dir,
#                     step,
#                     state,
#                     prefix='ckpt_',
#                     max_to_keep=None):
# def load_checkpoint(
#     checkpoint_path,
#     target=None,
#     prefix='ckpt_'):
# def load_latest_checkpoint(train_dir, target=None, prefix='ckpt_'):
# 
# 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/utils.py
import concurrent.futures
import functools
import json
import logging
import operator
import os.path
import time
import jax
import jax.numpy as jnp
import numpy as np
import pandas as pd
from typing import Any, Dict
from absl import logging as absl_logging
from clu import metric_writers
from flax import serialization
from flax import struct
from flax.training import checkpoints as flax_checkpoints
from init2winit import checkpoint
from init2winit.model_lib import model_utils
from tensorflow.io import gfile


  init_csv_path = os.path.join(train_dir, 'init_measurements.csv')
  init_json_path = os.path.join(train_dir, 'init_scalars.json')
  init_logger = MetricLogger(
      csv_path=init_csv_path,
      json_path=init_json_path,
      xm_work_unit=xm_work_unit)
  return metrics_logger, init_logger


def get_summary_tree(training_metrics_grabber):
  """Extracts desired training statistics from the grabber state.

  Currently this function will compute the scalar aggregate gradient variance
  for every weight matrix of the model. Future iterations of this function
  may depend on the metrics_grabber config.

  Args:
    training_metrics_grabber: TrainingMetricsGrabber object.

  Returns:
    A dict of different aggregate training statistics.
  """
  unreplicated_metrics_tree = jax.tree_map(lambda x: x[0],
                                           training_metrics_grabber.state)

  # Example key: Layer1/conv1/kernel/
  # NOTE: jax.tree_map does not work here, because tree_map will additionally
  # flatten the node state, while model_utils.flatten_dict will consider the
  # node object a leaf.

  flat_metrics = model_utils.flatten_dict(unreplicated_metrics_tree)