# Path: init2winit/dataset_lib/datasets.py
# def get_dataset(dataset_name):
#   """Maps dataset name to a dataset_builder."""
#   try:
#     return _ALL_DATASETS[dataset_name].getter
#   except KeyError:
#     raise ValueError('Unrecognized dataset: {}'.format(dataset_name))
# 
# Path: init2winit/dataset_lib/datasets.py
# def get_dataset_hparams(dataset_name):
#   """Maps dataset name to default_hps."""
#   try:
#     hparams = _ALL_DATASETS[dataset_name].hparams
#     # TODO(mbadura): Refactor to explicitly support different input specs
#     if 'input_shape' not in hparams or hparams.input_shape is None:
#       if 'input_edge_shape' in hparams and 'input_node_shape' in hparams:
#         hparams.input_shape = (hparams.input_node_shape,
#                                hparams.input_edge_shape)
#       elif dataset_name == 'lm1b':
#         max_len = max(hparams.max_target_length, hparams.max_eval_target_length)
#         hparams.input_shape = (max_len,)
#       elif dataset_name == 'translate_wmt':
#         max_len = max(hparams.max_target_length, hparams.max_eval_target_length,
#                       hparams.max_predict_length)
#         hparams.input_shape = [(max_len,), (max_len,)]
#       else:
#         raise ValueError(
#             'Undefined input shape for dataset: {}'.format(dataset_name))
#     return hparams
#   except KeyError:
#     raise ValueError('Unrecognized dataset: {}'.format(dataset_name))


import itertools
import jax.numpy as jnp
import jax.random
import numpy as np
import tensorflow_datasets as tfds
from absl import flags
from absl.testing import absltest
from absl.testing import parameterized
from init2winit.dataset_lib.datasets import get_dataset
from init2winit.dataset_lib.datasets import get_dataset_hparams
# 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.

"""Unit tests for datasets.py."""




FLAGS = flags.FLAGS


# NOTE: Imagenet will not satisfy data determinism, so we will not test for it.
small_image_ds = ['mnist', 'cifar10', 'fashion_mnist']
determinism_test_keys = [('test_{}'.format(m), m) for m in small_image_ds]


class DatasetTest(parameterized.TestCase):
  """Unit tests for datasets.py."""

  @parameterized.named_parameters(*determinism_test_keys)
  def test_determinism(self, ds):
    """Test that shuffle_rng and epoch correctly determine the order of data."""
    batch_size = 32
    eval_batch_size = 16
    np.random.seed(0)  # set the seed so the mock data is deterministic.

    # This will override the tfds.load(mnist) call to return 100 fake samples.
    with tfds.testing.mock_data(num_examples=128):

# Path: init2winit/dataset_lib/datasets.py
# def get_dataset(dataset_name):
#   """Maps dataset name to a dataset_builder."""
#   try:
#     return _ALL_DATASETS[dataset_name].getter
#   except KeyError:
#     raise ValueError('Unrecognized dataset: {}'.format(dataset_name))
# 
# Path: init2winit/dataset_lib/datasets.py
# def get_dataset_hparams(dataset_name):
#   """Maps dataset name to default_hps."""
#   try:
#     hparams = _ALL_DATASETS[dataset_name].hparams
#     # TODO(mbadura): Refactor to explicitly support different input specs
#     if 'input_shape' not in hparams or hparams.input_shape is None:
#       if 'input_edge_shape' in hparams and 'input_node_shape' in hparams:
#         hparams.input_shape = (hparams.input_node_shape,
#                                hparams.input_edge_shape)
#       elif dataset_name == 'lm1b':
#         max_len = max(hparams.max_target_length, hparams.max_eval_target_length)
#         hparams.input_shape = (max_len,)
#       elif dataset_name == 'translate_wmt':
#         max_len = max(hparams.max_target_length, hparams.max_eval_target_length,
#                       hparams.max_predict_length)
#         hparams.input_shape = [(max_len,), (max_len,)]
#       else:
#         raise ValueError(
#             'Undefined input shape for dataset: {}'.format(dataset_name))
#     return hparams
#   except KeyError:
#     raise ValueError('Unrecognized dataset: {}'.format(dataset_name))


# Path: init2winit/dataset_lib/test_datasets.py
import itertools
import jax.numpy as jnp
import jax.random
import numpy as np
import tensorflow_datasets as tfds
from absl import flags
from absl.testing import absltest
from absl.testing import parameterized
from init2winit.dataset_lib.datasets import get_dataset
from init2winit.dataset_lib.datasets import get_dataset_hparams

# 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.

"""Unit tests for datasets.py."""




FLAGS = flags.FLAGS


# NOTE: Imagenet will not satisfy data determinism, so we will not test for it.
small_image_ds = ['mnist', 'cifar10', 'fashion_mnist']
determinism_test_keys = [('test_{}'.format(m), m) for m in small_image_ds]


class DatasetTest(parameterized.TestCase):
  """Unit tests for datasets.py."""

  @parameterized.named_parameters(*determinism_test_keys)
  def test_determinism(self, ds):
    """Test that shuffle_rng and epoch correctly determine the order of data."""
    batch_size = 32
    eval_batch_size = 16
    np.random.seed(0)  # set the seed so the mock data is deterministic.

    # This will override the tfds.load(mnist) call to return 100 fake samples.
    with tfds.testing.mock_data(num_examples=128):

      dataset_builder = get_dataset(ds)