# 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
    # The flattened keys are tuples, which doesn't work with the serialization,
    # so to save this dict the keys need to be strings.
    sub_results[str(key)] = row

  return sub_results


def prepare_batches_gen(dataset, eval_config):
  """Returns a data iterator.

  The API for the data iterator will be
    for b in batches_gen():
      pass
  We yield the same "epoch" every time to the data iterator is called.

  Args:
    dataset: An init2winit.dataset_lib.Dataset object. This is ignored if
      eval_config['use_training_gen'] == False.
    eval_config: A dict specifying the parameters for the hessian eval.

  Returns:
    A data generator.
  """
  train_iter = itertools.islice(dataset.train_iterator_fn(), 0,
                                eval_config['num_batches'])
  batches = list(train_iter)
  init_rng = jax.random.PRNGKey(eval_config['rng_key'])
  init_rng = jax.random.fold_in(init_rng, jax.process_index())
  def training_batches_gen():
    for counter, batch in enumerate(batches):

# 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

    # The flattened keys are tuples, which doesn't work with the serialization,
    # so to save this dict the keys need to be strings.
    sub_results[str(key)] = row

  return sub_results


def prepare_batches_gen(dataset, eval_config):
  """Returns a data iterator.

  The API for the data iterator will be
    for b in batches_gen():
      pass
  We yield the same "epoch" every time to the data iterator is called.

  Args:
    dataset: An init2winit.dataset_lib.Dataset object. This is ignored if
      eval_config['use_training_gen'] == False.
    eval_config: A dict specifying the parameters for the hessian eval.

  Returns:
    A data generator.
  """
  train_iter = itertools.islice(dataset.train_iterator_fn(), 0,
                                eval_config['num_batches'])
  batches = list(train_iter)
  init_rng = jax.random.PRNGKey(eval_config['rng_key'])
  init_rng = jax.random.fold_in(init_rng, jax.process_index())
  def training_batches_gen():
    for counter, batch in enumerate(batches):

      batch = data_utils.shard(batch)