# 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
  input_dtype = tf.bfloat16
  shuffle_buffer_size = 16384

  train_ds = mlperf_input_pipeline.load_split(
      host_batch_size,
      dtype=input_dtype,
      split='train',
      rng=rng,
      shuffle_size=shuffle_buffer_size,
      preprocess_fn=_preprocess_fn)

  eval_train_ds = mlperf_input_pipeline.load_split(
      host_batch_size,
      dtype=input_dtype,
      split='eval_train',
      rng=rng,
      shuffle_size=shuffle_buffer_size,
      preprocess_fn=_preprocess_fn)

  eval_ds = mlperf_input_pipeline.load_split(
      eval_host_batch_size,
      dtype=input_dtype,
      split='validation',
      rng=rng,
      shuffle_size=shuffle_buffer_size,
      preprocess_fn=_preprocess_fn)

  # We cannot use tfds.as_numpy because this calls tensor.numpy() which does an
  # additional copy of the tensor, instead we call tensor._numpy() below.
  def train_iterator_fn():

# 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

  input_dtype = tf.bfloat16
  shuffle_buffer_size = 16384

  train_ds = mlperf_input_pipeline.load_split(
      host_batch_size,
      dtype=input_dtype,
      split='train',
      rng=rng,
      shuffle_size=shuffle_buffer_size,
      preprocess_fn=_preprocess_fn)

  eval_train_ds = mlperf_input_pipeline.load_split(
      host_batch_size,
      dtype=input_dtype,
      split='eval_train',
      rng=rng,
      shuffle_size=shuffle_buffer_size,
      preprocess_fn=_preprocess_fn)

  eval_ds = mlperf_input_pipeline.load_split(
      eval_host_batch_size,
      dtype=input_dtype,
      split='validation',
      rng=rng,
      shuffle_size=shuffle_buffer_size,
      preprocess_fn=_preprocess_fn)

  # We cannot use tfds.as_numpy because this calls tensor.numpy() which does an
  # additional copy of the tensor, instead we call tensor._numpy() below.
  def train_iterator_fn():

    return data_utils.iterator_as_numpy(iter(train_ds))