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


import os
import shutil
import tempfile
from absl.testing import absltest
from absl.testing import parameterized
from init2winit import checkpoint
from init2winit import utils
          num_workers=3,
          input_list_dict=[dict(i=k) for k in range(1, 10)],
          expected=list(range(1, 10)),
      ),
      dict(
          testcase_name='more workers than jobs, nonempty list',
          num_workers=20,
          input_list_dict=[dict(i=k) for k in range(1, 10)],
          expected=list(range(1, 10)),
      ),
  )
  def testRunInParallel(self, input_list_dict, num_workers, expected):
    """Test running successful fns in parallel, originally from mlbileschi."""
    actual = utils.run_in_parallel(_identity, input_list_dict, num_workers)
    self.assertEqual(actual, expected)

  def testRunInParallelOnFailingFn(self):
    """Test running failing fns in parallel, originally from mlbileschi."""
    with self.assertRaisesRegex(ValueError, 'I always fail.'):
      utils.run_in_parallel(_fn_that_always_fails, [dict(arg='hi')], 10)

  def testAppendPytree(self):
    """Test appending and loading pytrees."""
    pytrees = [{'a': i} for i in range(10)]
    pytree_path = os.path.join(self.test_dir, 'pytree.ckpt')
    logger = utils.MetricLogger(pytree_path=pytree_path)

    for pytree in pytrees:
      logger.append_pytree(pytree)


# 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/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/test_utils.py
import os
import shutil
import tempfile
from absl.testing import absltest
from absl.testing import parameterized
from init2winit import checkpoint
from init2winit import utils

          num_workers=3,
          input_list_dict=[dict(i=k) for k in range(1, 10)],
          expected=list(range(1, 10)),
      ),
      dict(
          testcase_name='more workers than jobs, nonempty list',
          num_workers=20,
          input_list_dict=[dict(i=k) for k in range(1, 10)],
          expected=list(range(1, 10)),
      ),
  )
  def testRunInParallel(self, input_list_dict, num_workers, expected):
    """Test running successful fns in parallel, originally from mlbileschi."""
    actual = utils.run_in_parallel(_identity, input_list_dict, num_workers)
    self.assertEqual(actual, expected)

  def testRunInParallelOnFailingFn(self):
    """Test running failing fns in parallel, originally from mlbileschi."""
    with self.assertRaisesRegex(ValueError, 'I always fail.'):
      utils.run_in_parallel(_fn_that_always_fails, [dict(arg='hi')], 10)

  def testAppendPytree(self):
    """Test appending and loading pytrees."""
    pytrees = [{'a': i} for i in range(10)]
    pytree_path = os.path.join(self.test_dir, 'pytree.ckpt')
    logger = utils.MetricLogger(pytree_path=pytree_path)

    for pytree in pytrees:
      logger.append_pytree(pytree)


    latest = checkpoint.load_latest_checkpoint(pytree_path, prefix='')