# 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/dataset_lib/image_preprocessing.py
# def _crop(key, image, hps):
# def crop(key, images, hps):
# def mixup_general(key,
#                   *args_to_mix,
#                   alpha=0.1,
#                   fold_in=None,
#                   n=2,
#                   **kwargs_to_mix):
#   def mix(batch):
#     def mul(a, b):  # B * BHWC -> B111 * BHWC
# def mixup(key, alpha, images, labels):
# def augment_cifar10(key, images, labels, hps):
# def identity_augment(key, images, labels):
# 
# 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):


import functools
import itertools
import jax.numpy as jnp
import jax.random
import numpy as np
import tensorflow.compat.v2 as tf
import tensorflow_datasets as tfds
from absl import logging
from init2winit.dataset_lib import data_utils
from init2winit.dataset_lib import image_preprocessing
from init2winit.dataset_lib.data_utils import Dataset
from ml_collections.config_dict import config_dict
  # number, so everytime we generate epoch 10 it yields the same pseudorandom
  # order.

  train_iterator_fn = functools.partial(
      data_utils.image_iterator,
      data_train,
      rescale=rescale,
      output_shape=output_shape,
      is_one_hot=is_one_hot,
      autoencoder=autoencoder,
      shuffle_rng=shuffle_rng,
      augment_fn=augment_fn)

  eval_train_epoch = functools.partial(
      _eval_batches,
      eval_train_inputs,
      eval_train_targets,
      per_host_eval_batch_size)
  valid_epoch = functools.partial(
      _eval_batches,
      valid_inputs,
      valid_targets,
      per_host_eval_batch_size,
      valid_example_keys=valid_example_keys)
  test_epoch = functools.partial(
      _eval_batches,
      test_inputs,
      test_targets,
      per_host_eval_batch_size)


# 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/dataset_lib/image_preprocessing.py
# def _crop(key, image, hps):
# def crop(key, images, hps):
# def mixup_general(key,
#                   *args_to_mix,
#                   alpha=0.1,
#                   fold_in=None,
#                   n=2,
#                   **kwargs_to_mix):
#   def mix(batch):
#     def mul(a, b):  # B * BHWC -> B111 * BHWC
# def mixup(key, alpha, images, labels):
# def augment_cifar10(key, images, labels, hps):
# def identity_augment(key, images, labels):
# 
# 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/dataset_lib/small_image_datasets.py
import functools
import itertools
import jax.numpy as jnp
import jax.random
import numpy as np
import tensorflow.compat.v2 as tf
import tensorflow_datasets as tfds
from absl import logging
from init2winit.dataset_lib import data_utils
from init2winit.dataset_lib import image_preprocessing
from init2winit.dataset_lib.data_utils import Dataset
from ml_collections.config_dict import config_dict

  # number, so everytime we generate epoch 10 it yields the same pseudorandom
  # order.

  train_iterator_fn = functools.partial(
      data_utils.image_iterator,
      data_train,
      rescale=rescale,
      output_shape=output_shape,
      is_one_hot=is_one_hot,
      autoencoder=autoencoder,
      shuffle_rng=shuffle_rng,
      augment_fn=augment_fn)

  eval_train_epoch = functools.partial(
      _eval_batches,
      eval_train_inputs,
      eval_train_targets,
      per_host_eval_batch_size)
  valid_epoch = functools.partial(
      _eval_batches,
      valid_inputs,
      valid_targets,
      per_host_eval_batch_size,
      valid_example_keys=valid_example_keys)
  test_epoch = functools.partial(
      _eval_batches,
      test_inputs,
      test_targets,
      per_host_eval_batch_size)


  return Dataset(train_iterator_fn, eval_train_epoch, valid_epoch, test_epoch)