# 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
  eval_train_targets = np.reshape(eval_train_targets, eval_train_output_shape)

  valid_inputs = jnp.array([])
  valid_targets = jnp.array([])
  valid_example_keys = jnp.array([])
  if data_valid:
    valid_data = next(eval_image_iterator(
        data_valid, include_example_keys=include_example_keys))
    valid_inputs = valid_data['inputs']
    valid_targets = valid_data['targets']
    if include_example_keys:
      valid_example_keys = valid_data['tfds_id']

  test_data = next(eval_image_iterator(data_test))
  test_inputs = jnp.array(test_data['inputs'].astype(np.float32))
  test_targets = test_data['targets']

  # Shard the validation and test data by host id.
  valid_inputs = _shard_by_host_id(valid_inputs)
  valid_targets = _shard_by_host_id(valid_targets)
  test_inputs = _shard_by_host_id(test_inputs)
  test_targets = _shard_by_host_id(test_targets)

  # TODO(gilmer): The simplest way to do this would be to directly yield from
  # tfds.as_numpy(). However we currently do not know how to properly handle
  # restarts with tfds. We'd like the train shuffle to depend on the epoch
  # number, so everytime we generate epoch 10 it yields the same pseudorandom
  # order.

  train_iterator_fn = functools.partial(

# 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

  eval_train_targets = np.reshape(eval_train_targets, eval_train_output_shape)

  valid_inputs = jnp.array([])
  valid_targets = jnp.array([])
  valid_example_keys = jnp.array([])
  if data_valid:
    valid_data = next(eval_image_iterator(
        data_valid, include_example_keys=include_example_keys))
    valid_inputs = valid_data['inputs']
    valid_targets = valid_data['targets']
    if include_example_keys:
      valid_example_keys = valid_data['tfds_id']

  test_data = next(eval_image_iterator(data_test))
  test_inputs = jnp.array(test_data['inputs'].astype(np.float32))
  test_targets = test_data['targets']

  # Shard the validation and test data by host id.
  valid_inputs = _shard_by_host_id(valid_inputs)
  valid_targets = _shard_by_host_id(valid_targets)
  test_inputs = _shard_by_host_id(test_inputs)
  test_targets = _shard_by_host_id(test_targets)

  # TODO(gilmer): The simplest way to do this would be to directly yield from
  # tfds.as_numpy(). However we currently do not know how to properly handle
  # restarts with tfds. We'd like the train shuffle to depend on the epoch
  # number, so everytime we generate epoch 10 it yields the same pseudorandom
  # order.

  train_iterator_fn = functools.partial(

      data_utils.image_iterator,