# 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
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""ImageNet input pipeline with MLPerf preprocessing."""




NUM_CLASSES = 1000


DEFAULT_HPARAMS = config_dict.ConfigDict(dict(
    input_shape=(224, 224, 3),
    output_shape=(NUM_CLASSES,),
    train_size=1281167,
    valid_size=50000))

METADATA = {
    'apply_one_hot_in_loss': False,
}


def transpose_and_normalize_image(image):
  mean = tf.constant([[mlperf_input_pipeline.MEAN_RGB]], dtype=image.dtype)
  stddev = tf.constant(

# 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

#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""ImageNet input pipeline with MLPerf preprocessing."""




NUM_CLASSES = 1000


DEFAULT_HPARAMS = config_dict.ConfigDict(dict(
    input_shape=(224, 224, 3),
    output_shape=(NUM_CLASSES,),
    train_size=1281167,
    valid_size=50000))

METADATA = {
    'apply_one_hot_in_loss': False,
}


def transpose_and_normalize_image(image):
  mean = tf.constant([[mlperf_input_pipeline.MEAN_RGB]], dtype=image.dtype)
  stddev = tf.constant(

      [[mlperf_input_pipeline.STDDEV_RGB]], dtype=image.dtype)