# 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 shutil
import tempfile
import numpy as np
from absl import flags
from absl.testing import absltest
from init2winit import utils
from jax import test_util as jtu
"""Unit tests for trainer.py.

"""



FLAGS = flags.FLAGS


class TrainingMetricsTest(jtu.JaxTestCase):
  """Tests the logged statistics from training_metrics_grabber."""

  def setUp(self):
    super(TrainingMetricsTest, self).setUp()
    self.test_dir = tempfile.mkdtemp()

  def tearDown(self):
    shutil.rmtree(self.test_dir)
    super(TrainingMetricsTest, self).tearDown()

  def test_grad_var(self):
    model_size = 10
    example_grads = [{
        'layer1': np.ones(model_size),
        'layer2': 3 * np.ones(model_size)
    }, {
        'layer1': 2 * np.ones(model_size),
        'layer2': np.ones(model_size)
    }]
    eval_config = {'ema_beta': 0.5}

# 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_training_metrics_grabber.py
import shutil
import tempfile
import numpy as np
from absl import flags
from absl.testing import absltest
from init2winit import utils
from jax import test_util as jtu

"""Unit tests for trainer.py.

"""



FLAGS = flags.FLAGS


class TrainingMetricsTest(jtu.JaxTestCase):
  """Tests the logged statistics from training_metrics_grabber."""

  def setUp(self):
    super(TrainingMetricsTest, self).setUp()
    self.test_dir = tempfile.mkdtemp()

  def tearDown(self):
    shutil.rmtree(self.test_dir)
    super(TrainingMetricsTest, self).tearDown()

  def test_grad_var(self):
    model_size = 10
    example_grads = [{
        'layer1': np.ones(model_size),
        'layer2': 3 * np.ones(model_size)
    }, {
        'layer1': 2 * np.ones(model_size),
        'layer2': np.ones(model_size)
    }]
    eval_config = {'ema_beta': 0.5}

    training_metrics_grabber = utils.TrainingMetricsGrabber.create(