# 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/init_lib/init_utils.py
# def initialize(flax_module,
#                initializer,
#                loss_fn,
#                input_shape,
#                output_shape,
#                hps,
#                rng,
#                metrics_logger,
#                fake_batch=None):
# 
# Path: init2winit/init_lib/initializers.py
# def noop(
#     loss_fn=None,
#     flax_module=None,
#     params=None,
#     hps=None,
#     input_shape=None,
#     output_shape=None,
#     rng_key=None,
#     metrics_logger=None,
# ):
# def get_initializer(initializer_name):
# def get_initializer_hparams(initializer_name):
# 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),
# }
# 
# Path: init2winit/optimizer_lib/optimizers.py
# def sgd(learning_rate, weight_decay, momentum=None, nesterov=False):
# def get_optimizer(hps, model=None):
# def _wrap_update_fn(opt_name, opt_update):
#   def update_fn(grads, optimizer_state, params, batch=None, batch_stats=None):
# 
# Path: init2winit/trainer_lib/trainer.py
# _GRAD_CLIP_EPS = 1e-6
# def evaluate(
#     params,
#     batch_stats,
#     batch_iter,
#     evaluate_batch_pmapped):
# def _inject_learning_rate(optimizer_state, lr):
# def update(
#     optimizer_state,
#     params,
#     batch_stats,
#     batch,
#     step,
#     lr,
#     rng,
#     local_device_index,
#     training_metrics_grabber,
#     training_cost,
#     grad_clip,
#     optimizer_update_fn):
#   def opt_cost(params):
# def _merge_and_apply_prefix(d1, d2, prefix):
# def eval_metrics(params, batch_stats, dataset, eval_num_batches,
#                  eval_train_num_batches, evaluate_batch_pmapped):
# def train(train_dir,
#           model,
#           dataset_builder,
#           initializer,
#           num_train_steps,
#           hps,
#           rng,
#           eval_batch_size,
#           eval_num_batches,
#           eval_train_num_batches,
#           eval_frequency,
#           checkpoint_steps,
#           early_stopping_target_name=None,
#           early_stopping_target_value=None,
#           early_stopping_mode=None,
#           eval_steps=None,
#           metrics_logger=None,
#           init_logger=None,
#           training_metrics_config=None,
#           callback_configs=None):
#   def get_step_frequency(cur_step):
# 
# Path: init2winit/trainer_lib/trainer_utils.py
# def log_epoch_report(report, metrics_logger):
# def maybe_log_training_metrics(training_metrics_grabber, metrics_logger):
# def maybe_sync_batchnorm_stats(batch_stats):
# def should_eval(global_step, eval_frequency, eval_steps):
# def check_for_early_stopping(
#     early_stopping_target_name,
#     early_stopping_target_value,
#     early_stopping_mode,
#     eval_report):


import json
import os
import utils as utils  # local file import
import jax
import optax
from absl import logging
from flax import jax_utils
from init2winit import checkpoint
from init2winit.hessian import hessian_eval
from init2winit.init_lib import init_utils
from init2winit.init_lib import initializers
from init2winit.optimizer_lib import optimizers
from init2winit.trainer_lib import trainer
from init2winit.trainer_lib import trainer_utils
from tensorflow.io import gfile

  if jax.process_index() == 0:
    logging.info('Starting eval!')
    logging.info('Number of hosts: %d', jax.process_count())

  hessian_evaluator = hessian_eval.CurvatureEvaluator(
      params,
      hessian_eval_config,
      dataset=dataset,
      loss=batch_loss)
  if min_global_step is None:
    suffix = ''
  else:
    suffix = '{}_{}'.format(min_global_step, max_global_step)
  pytree_path = os.path.join(checkpoint_dir, hessian_eval_config['name'],
                             suffix)
  logger = utils.MetricLogger(pytree_path=pytree_path)
  for checkpoint_path, step in iterate_checkpoints(checkpoint_dir,
                                                   min_global_step,
                                                   max_global_step):
    unreplicated_checkpoint_state = dict(
        params=unreplicated_params,
        optimizer_state=unreplicated_optimizer_state,
        batch_stats=unreplicated_batch_stats,
        global_step=0,
        preemption_count=0,
        sum_train_cost=0.0)
    ckpt = checkpoint.load_checkpoint(
        checkpoint_path,
        target=unreplicated_checkpoint_state)

# 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/init_lib/init_utils.py
# def initialize(flax_module,
#                initializer,
#                loss_fn,
#                input_shape,
#                output_shape,
#                hps,
#                rng,
#                metrics_logger,
#                fake_batch=None):
# 
# Path: init2winit/init_lib/initializers.py
# def noop(
#     loss_fn=None,
#     flax_module=None,
#     params=None,
#     hps=None,
#     input_shape=None,
#     output_shape=None,
#     rng_key=None,
#     metrics_logger=None,
# ):
# def get_initializer(initializer_name):
# def get_initializer_hparams(initializer_name):
# 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),
# }
# 
# Path: init2winit/optimizer_lib/optimizers.py
# def sgd(learning_rate, weight_decay, momentum=None, nesterov=False):
# def get_optimizer(hps, model=None):
# def _wrap_update_fn(opt_name, opt_update):
#   def update_fn(grads, optimizer_state, params, batch=None, batch_stats=None):
# 
# Path: init2winit/trainer_lib/trainer.py
# _GRAD_CLIP_EPS = 1e-6
# def evaluate(
#     params,
#     batch_stats,
#     batch_iter,
#     evaluate_batch_pmapped):
# def _inject_learning_rate(optimizer_state, lr):
# def update(
#     optimizer_state,
#     params,
#     batch_stats,
#     batch,
#     step,
#     lr,
#     rng,
#     local_device_index,
#     training_metrics_grabber,
#     training_cost,
#     grad_clip,
#     optimizer_update_fn):
#   def opt_cost(params):
# def _merge_and_apply_prefix(d1, d2, prefix):
# def eval_metrics(params, batch_stats, dataset, eval_num_batches,
#                  eval_train_num_batches, evaluate_batch_pmapped):
# def train(train_dir,
#           model,
#           dataset_builder,
#           initializer,
#           num_train_steps,
#           hps,
#           rng,
#           eval_batch_size,
#           eval_num_batches,
#           eval_train_num_batches,
#           eval_frequency,
#           checkpoint_steps,
#           early_stopping_target_name=None,
#           early_stopping_target_value=None,
#           early_stopping_mode=None,
#           eval_steps=None,
#           metrics_logger=None,
#           init_logger=None,
#           training_metrics_config=None,
#           callback_configs=None):
#   def get_step_frequency(cur_step):
# 
# Path: init2winit/trainer_lib/trainer_utils.py
# def log_epoch_report(report, metrics_logger):
# def maybe_log_training_metrics(training_metrics_grabber, metrics_logger):
# def maybe_sync_batchnorm_stats(batch_stats):
# def should_eval(global_step, eval_frequency, eval_steps):
# def check_for_early_stopping(
#     early_stopping_target_name,
#     early_stopping_target_value,
#     early_stopping_mode,
#     eval_report):


# Path: hessian/run_lanczos.py
import json
import os
import utils as utils  # local file import
import jax
import optax
from absl import logging
from flax import jax_utils
from init2winit import checkpoint
from init2winit.hessian import hessian_eval
from init2winit.init_lib import init_utils
from init2winit.init_lib import initializers
from init2winit.optimizer_lib import optimizers
from init2winit.trainer_lib import trainer
from init2winit.trainer_lib import trainer_utils
from tensorflow.io import gfile


  if jax.process_index() == 0:
    logging.info('Starting eval!')
    logging.info('Number of hosts: %d', jax.process_count())

  hessian_evaluator = hessian_eval.CurvatureEvaluator(
      params,
      hessian_eval_config,
      dataset=dataset,
      loss=batch_loss)
  if min_global_step is None:
    suffix = ''
  else:
    suffix = '{}_{}'.format(min_global_step, max_global_step)
  pytree_path = os.path.join(checkpoint_dir, hessian_eval_config['name'],
                             suffix)
  logger = utils.MetricLogger(pytree_path=pytree_path)
  for checkpoint_path, step in iterate_checkpoints(checkpoint_dir,
                                                   min_global_step,
                                                   max_global_step):
    unreplicated_checkpoint_state = dict(
        params=unreplicated_params,
        optimizer_state=unreplicated_optimizer_state,
        batch_stats=unreplicated_batch_stats,
        global_step=0,
        preemption_count=0,
        sum_train_cost=0.0)
    ckpt = checkpoint.load_checkpoint(
        checkpoint_path,
        target=unreplicated_checkpoint_state)

    results, _ = checkpoint.replicate_checkpoint(