# 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):


from init2winit.dataset_lib import data_utils
from ml_collections.config_dict import config_dict
import jax
import jax.numpy as jnp
import numpy as np
      per_host_eval_batch_size, hps.input_shape, hps.output_shape[0])

  def train_iterator_fn():
    while True:
      yield fake_train_batch

  def valid_epoch(epoch, num_batches=None):
    del num_batches
    del epoch
    # Note that we do // beacuse we do not support partial batching for the fake
    # dataset.
    for _ in range(hps.valid_size // eval_batch_size):
      yield fake_test_batch

  # pylint: disable=unreachable
  def eval_train_epoch(*args, **kwargs):
    del args
    del kwargs
    return
    yield  # This yield is needed to make this a valid (null) iterator.
  # pylint: enable=unreachable
  # 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/fake_dataset.py
from init2winit.dataset_lib import data_utils
from ml_collections.config_dict import config_dict
import jax
import jax.numpy as jnp
import numpy as np

      per_host_eval_batch_size, hps.input_shape, hps.output_shape[0])

  def train_iterator_fn():
    while True:
      yield fake_train_batch

  def valid_epoch(epoch, num_batches=None):
    del num_batches
    del epoch
    # Note that we do // beacuse we do not support partial batching for the fake
    # dataset.
    for _ in range(hps.valid_size // eval_batch_size):
      yield fake_test_batch

  # pylint: disable=unreachable
  def eval_train_epoch(*args, **kwargs):
    del args
    del kwargs
    return
    yield  # This yield is needed to make this a valid (null) iterator.
  # pylint: enable=unreachable
  # 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(