# 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
      with gfile.GFile(self._csv_path) as csv_file:
        measurements = pd.read_csv(csv_file)
      measurements = measurements.append([metrics])
    else:
      measurements = pd.DataFrame([metrics], columns=sorted(metrics.keys()))
    # TODO(gdahl,gilmer): Should this be an atomic file?
    with gfile.GFile(self._csv_path, 'w') as csv_file:
      measurements.to_csv(csv_file, index=False)
    if self._xm_work_unit:
      for name, value in metrics.items():
        if name not in self._measurements:
          self._measurements[name] = self._xm_work_unit.get_measurement_series(
              label=name)
        self._measurements[name].create_measurement(
            objective_value=value, step=metrics['global_step'])

    if self._tb_metric_writer:
      self._tb_metric_writer.write_scalars(
          step=int(metrics['global_step']), scalars=metrics)
      # This gives a 1-2% slowdown in steps_per_sec on cifar-10 with batch
      # size 512. We could only flush at the end of training to optimize this.
      self._tb_metric_writer.flush()

  def write_pytree(self, pytree):
    """Record a serializable pytree to disk, overwriting any previous state.

    Args:
      pytree: Any serializable pytree.
    """
    state = dict(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/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

      with gfile.GFile(self._csv_path) as csv_file:
        measurements = pd.read_csv(csv_file)
      measurements = measurements.append([metrics])
    else:
      measurements = pd.DataFrame([metrics], columns=sorted(metrics.keys()))
    # TODO(gdahl,gilmer): Should this be an atomic file?
    with gfile.GFile(self._csv_path, 'w') as csv_file:
      measurements.to_csv(csv_file, index=False)
    if self._xm_work_unit:
      for name, value in metrics.items():
        if name not in self._measurements:
          self._measurements[name] = self._xm_work_unit.get_measurement_series(
              label=name)
        self._measurements[name].create_measurement(
            objective_value=value, step=metrics['global_step'])

    if self._tb_metric_writer:
      self._tb_metric_writer.write_scalars(
          step=int(metrics['global_step']), scalars=metrics)
      # This gives a 1-2% slowdown in steps_per_sec on cifar-10 with batch
      # size 512. We could only flush at the end of training to optimize this.
      self._tb_metric_writer.flush()

  def write_pytree(self, pytree):
    """Record a serializable pytree to disk, overwriting any previous state.

    Args:
      pytree: Any serializable pytree.
    """
    state = dict(pytree=pytree)

    checkpoint.save_checkpoint(