# 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/mlperf_input_pipeline.py
# EVAL_IMAGES = 50000
# CROP_PADDING = 32
# MEAN_RGB = [0.485 * 255, 0.456 * 255, 0.406 * 255]
# STDDEV_RGB = [0.229 * 255, 0.224 * 255, 0.225 * 255]
# def load_split(batch_size,
#                split,
#                dtype,
#                rng,
#                data_dir=None,
#                image_size=224,
#                shuffle_size=16384,
#                preprocess_fn=None):
#   def dataset_parser(*value):


import itertools
import jax
import tensorflow.compat.v2 as tf
from init2winit.dataset_lib import data_utils
from init2winit.dataset_lib import mlperf_input_pipeline
from ml_collections.config_dict import config_dict
  # additional copy of the tensor, instead we call tensor._numpy() below.
  def train_iterator_fn():
    return data_utils.iterator_as_numpy(iter(train_ds))

  def eval_train_epoch(num_batches=None):
    if num_batches is None:
      num_batches = 0
    eval_train_iter = iter(eval_train_ds)
    np_iter = data_utils.iterator_as_numpy(
        itertools.islice(eval_train_iter, num_batches))
    for batch in np_iter:
      yield data_utils.maybe_pad_batch(batch, eval_host_batch_size)

  def valid_epoch(num_batches=None):
    if num_batches is None:
      num_batches = max_eval_steps
    valid_iter = iter(eval_ds)
    np_iter = data_utils.iterator_as_numpy(
        itertools.islice(valid_iter, num_batches))
    for batch in np_iter:
      yield data_utils.maybe_pad_batch(batch, eval_host_batch_size)

  # pylint: disable=unreachable
  def test_epoch(*args, **kwargs):
    del args
    del kwargs
    return
    yield  # This yield is needed to make this a valid (null) iterator.
  # pylint: enable=unreachable


# 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/mlperf_input_pipeline.py
# EVAL_IMAGES = 50000
# CROP_PADDING = 32
# MEAN_RGB = [0.485 * 255, 0.456 * 255, 0.406 * 255]
# STDDEV_RGB = [0.229 * 255, 0.224 * 255, 0.225 * 255]
# def load_split(batch_size,
#                split,
#                dtype,
#                rng,
#                data_dir=None,
#                image_size=224,
#                shuffle_size=16384,
#                preprocess_fn=None):
#   def dataset_parser(*value):


# Path: init2winit/dataset_lib/mlperf_imagenet_dataset.py
import itertools
import jax
import tensorflow.compat.v2 as tf
from init2winit.dataset_lib import data_utils
from init2winit.dataset_lib import mlperf_input_pipeline
from ml_collections.config_dict import config_dict

  # additional copy of the tensor, instead we call tensor._numpy() below.
  def train_iterator_fn():
    return data_utils.iterator_as_numpy(iter(train_ds))

  def eval_train_epoch(num_batches=None):
    if num_batches is None:
      num_batches = 0
    eval_train_iter = iter(eval_train_ds)
    np_iter = data_utils.iterator_as_numpy(
        itertools.islice(eval_train_iter, num_batches))
    for batch in np_iter:
      yield data_utils.maybe_pad_batch(batch, eval_host_batch_size)

  def valid_epoch(num_batches=None):
    if num_batches is None:
      num_batches = max_eval_steps
    valid_iter = iter(eval_ds)
    np_iter = data_utils.iterator_as_numpy(
        itertools.islice(valid_iter, num_batches))
    for batch in np_iter:
      yield data_utils.maybe_pad_batch(batch, eval_host_batch_size)

  # pylint: disable=unreachable
  def test_epoch(*args, **kwargs):
    del args
    del kwargs
    return
    yield  # This yield is needed to make this a valid (null) iterator.
  # pylint: enable=unreachable


  return data_utils.Dataset(