# 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.data_utils import Dataset
from ml_collections.config_dict import config_dict
import jax.random
import numpy as np
      yield {
          'inputs':
              train_rng.normal(
                  size=(per_host_batch_size, *hps.input_shape))
      }

  def eval_train_epoch(num_batches):
    for _ in range(num_batches):
      yield {
          'inputs':
              eval_rng.normal(
                  size=(per_host_batch_size, *hps.input_shape))
      }

  # pylint: disable=unreachable
  def valid_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/nqm_noise.py
from init2winit.dataset_lib.data_utils import Dataset
from ml_collections.config_dict import config_dict
import jax.random
import numpy as np

      yield {
          'inputs':
              train_rng.normal(
                  size=(per_host_batch_size, *hps.input_shape))
      }

  def eval_train_epoch(num_batches):
    for _ in range(num_batches):
      yield {
          'inputs':
              eval_rng.normal(
                  size=(per_host_batch_size, *hps.input_shape))
      }

  # pylint: disable=unreachable
  def valid_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 Dataset(train_iterator_fn, eval_train_epoch, valid_epoch, test_epoch)