# Path: init2winit/base_callback.py
# class BaseCallBack:
#   def __init__(self, model, optimizer, batch_stats, dataset, hps,
#                callback_config, train_dir, rng):
#   def run_eval(self, optimizer, batch_stats, global_step):
# 
# 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 jax
from init2winit import base_callback
from init2winit import utils
from init2winit.hessian import hessian_eval
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Callback which runs the hessian eval."""




def set_up_hessian_eval(model, flax_module, batch_stats, dataset,
                        checkpoint_dir, hessian_eval_config):
  """Builds the CurvatureEvaluator object."""

  # First copy then unreplicate batch_stats. Note batch_stats doesn't affect the
  # forward pass in the hessian eval because we always run the model in training
  # However, we need to provide batch_stats for the model.training_cost API.
  # The copy is needed b/c the trainer will modify underlying arrays.
  batch_stats = jax.tree_map(lambda x: x[:][0], batch_stats)
  def batch_loss(module, batch_rng):
    batch, rng = batch_rng
    return model.training_cost(
        module, batch, batch_stats=batch_stats, dropout_rng=rng)[0]

  pytree_path = os.path.join(checkpoint_dir, hessian_eval_config['name'])

# Path: init2winit/base_callback.py
# class BaseCallBack:
#   def __init__(self, model, optimizer, batch_stats, dataset, hps,
#                callback_config, train_dir, rng):
#   def run_eval(self, optimizer, batch_stats, global_step):
# 
# 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: hessian/hessian_callback.py
import os
import jax
from init2winit import base_callback
from init2winit import utils
from init2winit.hessian import hessian_eval

# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Callback which runs the hessian eval."""




def set_up_hessian_eval(model, flax_module, batch_stats, dataset,
                        checkpoint_dir, hessian_eval_config):
  """Builds the CurvatureEvaluator object."""

  # First copy then unreplicate batch_stats. Note batch_stats doesn't affect the
  # forward pass in the hessian eval because we always run the model in training
  # However, we need to provide batch_stats for the model.training_cost API.
  # The copy is needed b/c the trainer will modify underlying arrays.
  batch_stats = jax.tree_map(lambda x: x[:][0], batch_stats)
  def batch_loss(module, batch_rng):
    batch, rng = batch_rng
    return model.training_cost(
        module, batch, batch_stats=batch_stats, dropout_rng=rng)[0]

  pytree_path = os.path.join(checkpoint_dir, hessian_eval_config['name'])

  logger = utils.MetricLogger(pytree_path=pytree_path)