# Path: init2winit/dataset_lib/autoaugment.py
# _MAX_LEVEL = 10.
# NAME_TO_FUNC = {
#     'AutoContrast': autocontrast,
#     'Equalize': equalize,
#     'Invert': invert,
#     'Rotate': rotate,
#     'Posterize': posterize,
#     'Solarize': solarize,
#     'SolarizeAdd': solarize_add,
#     'Color': color,
#     'Contrast': contrast,
#     'Brightness': brightness,
#     'Sharpness': sharpness,
#     'ShearX': shear_x,
#     'ShearY': shear_y,
#     'TranslateX': translate_x,
#     'TranslateY': translate_y,
#     'Cutout': cutout,
# }
# def policy_v0():
# def policy_vtest():
# def blend(image1, image2, factor):
# def cutout(image, pad_size, replace=0):
# def solarize(image, threshold=128):
# def solarize_add(image, addition=0, threshold=128):
# def color(image, factor):
# def contrast(image, factor):
# def brightness(image, factor):
# def posterize(image, bits):
# def rotate(image, degrees, replace):
# def translate_x(image, pixels, replace):
# def translate_y(image, pixels, replace):
# def shear_x(image, level, replace):
# def shear_y(image, level, replace):
# def autocontrast(image):
#   def scale_channel(image):
#     def scale_values(im):
# def sharpness(image, factor):
# def equalize(image):
#   def scale_channel(im, c):
#     def build_lut(histo, step):
# def invert(image):
# def wrap(image):
# def unwrap(image, replace):
# def _randomly_negate_tensor(tensor):
# def _rotate_level_to_arg(level):
# def _shrink_level_to_arg(level):
# def _enhance_level_to_arg(level):
# def _shear_level_to_arg(level):
# def _translate_level_to_arg(level, translate_const):
# def level_to_arg(hparams):
# def _parse_policy_info(name, prob, level, replace_value, augmentation_hparams):
# def _apply_func_with_prob(func, image, args, prob):
# def select_and_apply_random_policy(policies, image):
# def build_and_apply_nas_policy(policies, image,
#                                augmentation_hparams):
#     def make_final_policy(tf_policy_):
#       def final_policy(image_):
# def distort_image_with_autoaugment(image, augmentation_name):
# def distort_image_with_randaugment(image, num_layers, magnitude, key):
# 
# 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):


from collections import abc
from absl import logging
from init2winit.dataset_lib import autoaugment
from init2winit.dataset_lib import data_utils
from init2winit.dataset_lib import image_preprocessing
from ml_collections.config_dict import config_dict
import itertools
import jax
import jax.numpy as jnp
import numpy as np
import tensorflow.compat.v2 as tf
import tensorflow_datasets as tfds
            targets,
            alpha=hps.mixup.alpha,
            n=2)
      yield {
          'inputs': image,
          'targets': targets,
          'weights': np.ones(per_host_batch_size, dtype=image.dtype)
      }

  def eval_train_epoch(num_batches=None):
    # This uses per_host_batch_size and not per_host_eval_batch_size.
    eval_train_iter = iter(eval_train_ds)
    for batch in itertools.islice(eval_train_iter, num_batches):
      batch_dict = {
          'inputs': batch['image'],
          'targets': np.eye(num_classes)[batch['label']],
      }
      if hps.get('include_example_keys'):
        batch_dict['example_key'] = batch['example_key']
      yield data_utils.maybe_pad_batch(batch_dict, per_host_eval_batch_size)

  def valid_epoch(num_batches=None):
    valid_iter = iter(eval_ds)
    for batch in itertools.islice(valid_iter, num_batches):
      batch_dict = {
          'inputs': batch['image'],
          'targets': np.eye(num_classes)[batch['label']],
      }
      if hps.get('include_example_keys'):
        batch_dict['example_key'] = batch['example_key']

# Path: init2winit/dataset_lib/autoaugment.py
# _MAX_LEVEL = 10.
# NAME_TO_FUNC = {
#     'AutoContrast': autocontrast,
#     'Equalize': equalize,
#     'Invert': invert,
#     'Rotate': rotate,
#     'Posterize': posterize,
#     'Solarize': solarize,
#     'SolarizeAdd': solarize_add,
#     'Color': color,
#     'Contrast': contrast,
#     'Brightness': brightness,
#     'Sharpness': sharpness,
#     'ShearX': shear_x,
#     'ShearY': shear_y,
#     'TranslateX': translate_x,
#     'TranslateY': translate_y,
#     'Cutout': cutout,
# }
# def policy_v0():
# def policy_vtest():
# def blend(image1, image2, factor):
# def cutout(image, pad_size, replace=0):
# def solarize(image, threshold=128):
# def solarize_add(image, addition=0, threshold=128):
# def color(image, factor):
# def contrast(image, factor):
# def brightness(image, factor):
# def posterize(image, bits):
# def rotate(image, degrees, replace):
# def translate_x(image, pixels, replace):
# def translate_y(image, pixels, replace):
# def shear_x(image, level, replace):
# def shear_y(image, level, replace):
# def autocontrast(image):
#   def scale_channel(image):
#     def scale_values(im):
# def sharpness(image, factor):
# def equalize(image):
#   def scale_channel(im, c):
#     def build_lut(histo, step):
# def invert(image):
# def wrap(image):
# def unwrap(image, replace):
# def _randomly_negate_tensor(tensor):
# def _rotate_level_to_arg(level):
# def _shrink_level_to_arg(level):
# def _enhance_level_to_arg(level):
# def _shear_level_to_arg(level):
# def _translate_level_to_arg(level, translate_const):
# def level_to_arg(hparams):
# def _parse_policy_info(name, prob, level, replace_value, augmentation_hparams):
# def _apply_func_with_prob(func, image, args, prob):
# def select_and_apply_random_policy(policies, image):
# def build_and_apply_nas_policy(policies, image,
#                                augmentation_hparams):
#     def make_final_policy(tf_policy_):
#       def final_policy(image_):
# def distort_image_with_autoaugment(image, augmentation_name):
# def distort_image_with_randaugment(image, num_layers, magnitude, key):
# 
# 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/imagenet_dataset.py
from collections import abc
from absl import logging
from init2winit.dataset_lib import autoaugment
from init2winit.dataset_lib import data_utils
from init2winit.dataset_lib import image_preprocessing
from ml_collections.config_dict import config_dict
import itertools
import jax
import jax.numpy as jnp
import numpy as np
import tensorflow.compat.v2 as tf
import tensorflow_datasets as tfds

            targets,
            alpha=hps.mixup.alpha,
            n=2)
      yield {
          'inputs': image,
          'targets': targets,
          'weights': np.ones(per_host_batch_size, dtype=image.dtype)
      }

  def eval_train_epoch(num_batches=None):
    # This uses per_host_batch_size and not per_host_eval_batch_size.
    eval_train_iter = iter(eval_train_ds)
    for batch in itertools.islice(eval_train_iter, num_batches):
      batch_dict = {
          'inputs': batch['image'],
          'targets': np.eye(num_classes)[batch['label']],
      }
      if hps.get('include_example_keys'):
        batch_dict['example_key'] = batch['example_key']
      yield data_utils.maybe_pad_batch(batch_dict, per_host_eval_batch_size)

  def valid_epoch(num_batches=None):
    valid_iter = iter(eval_ds)
    for batch in itertools.islice(valid_iter, num_batches):
      batch_dict = {
          'inputs': batch['image'],
          'targets': np.eye(num_classes)[batch['label']],
      }
      if hps.get('include_example_keys'):
        batch_dict['example_key'] = batch['example_key']

      yield data_utils.maybe_pad_batch(batch_dict, per_host_eval_batch_size)