# 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
    crop_num_pixels=4,
    use_mixup=False,
    train_size=73257 - 7000,
    valid_size=7000,
    test_size=26032,
    input_shape=(32, 32, 3),
    output_shape=(10,)))
SVHN_NO_EXTRA_METADATA = {
    'apply_one_hot_in_loss': False,
}


def _eval_batches(images,
                  labels,
                  per_host_batch_size,
                  num_batches=None,
                  valid_example_keys=None):
  """Produce a stream of batches for a single evaluation epoch."""
  for idx in itertools.islice(
      range(0, images.shape[0], per_host_batch_size), num_batches):
    inputs = jnp.array(images[idx:idx + per_host_batch_size])
    targets = jnp.array(labels[idx:idx + per_host_batch_size])
    data_dict = {
        'inputs': inputs,
        'targets': targets,
        'weights': jnp.ones(per_host_batch_size, dtype=inputs.dtype),
    }
    if valid_example_keys is not None:
      data_dict['example_key'] = valid_example_keys[idx:idx +
                                                    per_host_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

    crop_num_pixels=4,
    use_mixup=False,
    train_size=73257 - 7000,
    valid_size=7000,
    test_size=26032,
    input_shape=(32, 32, 3),
    output_shape=(10,)))
SVHN_NO_EXTRA_METADATA = {
    'apply_one_hot_in_loss': False,
}


def _eval_batches(images,
                  labels,
                  per_host_batch_size,
                  num_batches=None,
                  valid_example_keys=None):
  """Produce a stream of batches for a single evaluation epoch."""
  for idx in itertools.islice(
      range(0, images.shape[0], per_host_batch_size), num_batches):
    inputs = jnp.array(images[idx:idx + per_host_batch_size])
    targets = jnp.array(labels[idx:idx + per_host_batch_size])
    data_dict = {
        'inputs': inputs,
        'targets': targets,
        'weights': jnp.ones(per_host_batch_size, dtype=inputs.dtype),
    }
    if valid_example_keys is not None:
      data_dict['example_key'] = valid_example_keys[idx:idx +
                                                    per_host_batch_size]

    yield data_utils.maybe_pad_batch(data_dict, per_host_batch_size)