# Path: init2winit/dataset_lib/data_utils.py
# def iterator_as_numpy(iterator):
# def image_iterator(data,
#                    rescale,
#                    output_shape,
#                    is_one_hot,
#                    autoencoder,
#                    shuffle_rng=None,
#                    augment_fn=None,
#                    include_example_keys=False):
# def maybe_pad_batch(batch,
#                     desired_batch_size,
#                     data_format=None,
#                     mask_key=None):
#   def zero_pad(ar, pad_axis):
# def shard(batch, n_devices=None):
#   def _shard_array(array):
# def tf_to_numpy(tfds_data):
# 
# Path: init2winit/model_lib/partition_tree.py
# def outer_key(x):
# def create_partition_flat_params_fn(key_map):
#   def partition_flat_params(flat_params):
# def get_param_partition_fn(name):


import functools
import itertools
import flax
import jax
import jax.numpy as jnp
import numpy as np
import optax
import spectral_density.hessian_computation as hessian_computation
import spectral_density.lanczos as lanczos
from absl import logging
from flax import jax_utils
from flax.core.frozen_dict import unfreeze
from init2winit.dataset_lib import data_utils
from init2winit.model_lib import partition_tree
from spectral_density.hessian_computation import ravel_pytree
    count = 0.0
    full_grad = compiled_tree_zeros_like(params)
    for batch in self.batches_gen():
      # Already pmeaned minibatch gradients
      avg_grad = self.grad_loss(params, batch)
      if count < self.eval_config['num_eval_draws']:
        count += 1.0
        full_grad = compiled_tree_sum(full_grad, avg_grad)
    full_grad = jax.tree_map(lambda x: x / count, full_grad)
    full_grad = _unreplicate(full_grad)
    return full_grad

  def evaluate_spectrum(self, params, step):
    """Estimate the eignespectrum of H and C."""
    # Number of upper and lower eigenvectors to be approximated.
    num_evs = self.eval_config['num_eigens']
    hess_evecs = []
    cov_evecs = []

    if 2 * num_evs >= self.eval_config['num_lanczos_steps']:
      raise ValueError('Too many eigenvectors requested!')
    hvp_cl = lambda v: self.hvp_fn(params, v)
    gvp_cl = lambda v: self.gvp_fn(params, v)

    key = jax.random.PRNGKey(0)

    row = {'step': step}

    # Do we want .get here to keep old configs working?
    if self.eval_config.get('block_hessian'):

# Path: init2winit/dataset_lib/data_utils.py
# def iterator_as_numpy(iterator):
# def image_iterator(data,
#                    rescale,
#                    output_shape,
#                    is_one_hot,
#                    autoencoder,
#                    shuffle_rng=None,
#                    augment_fn=None,
#                    include_example_keys=False):
# def maybe_pad_batch(batch,
#                     desired_batch_size,
#                     data_format=None,
#                     mask_key=None):
#   def zero_pad(ar, pad_axis):
# def shard(batch, n_devices=None):
#   def _shard_array(array):
# def tf_to_numpy(tfds_data):
# 
# Path: init2winit/model_lib/partition_tree.py
# def outer_key(x):
# def create_partition_flat_params_fn(key_map):
#   def partition_flat_params(flat_params):
# def get_param_partition_fn(name):


# Path: hessian/hessian_eval.py
import functools
import itertools
import flax
import jax
import jax.numpy as jnp
import numpy as np
import optax
import spectral_density.hessian_computation as hessian_computation
import spectral_density.lanczos as lanczos
from absl import logging
from flax import jax_utils
from flax.core.frozen_dict import unfreeze
from init2winit.dataset_lib import data_utils
from init2winit.model_lib import partition_tree
from spectral_density.hessian_computation import ravel_pytree

    count = 0.0
    full_grad = compiled_tree_zeros_like(params)
    for batch in self.batches_gen():
      # Already pmeaned minibatch gradients
      avg_grad = self.grad_loss(params, batch)
      if count < self.eval_config['num_eval_draws']:
        count += 1.0
        full_grad = compiled_tree_sum(full_grad, avg_grad)
    full_grad = jax.tree_map(lambda x: x / count, full_grad)
    full_grad = _unreplicate(full_grad)
    return full_grad

  def evaluate_spectrum(self, params, step):
    """Estimate the eignespectrum of H and C."""
    # Number of upper and lower eigenvectors to be approximated.
    num_evs = self.eval_config['num_eigens']
    hess_evecs = []
    cov_evecs = []

    if 2 * num_evs >= self.eval_config['num_lanczos_steps']:
      raise ValueError('Too many eigenvectors requested!')
    hvp_cl = lambda v: self.hvp_fn(params, v)
    gvp_cl = lambda v: self.gvp_fn(params, v)

    key = jax.random.PRNGKey(0)

    row = {'step': step}

    # Do we want .get here to keep old configs working?
    if self.eval_config.get('block_hessian'):

      param_partition_fn = partition_tree.get_param_partition_fn(