# 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
    self.assertTrue(
        jnp.array_equal(saved_batch['inputs'],
                        saved_batch_same_epoch['inputs']))
    self.assertTrue(
        jnp.array_equal(saved_batch['targets'],
                        saved_batch_same_epoch['targets']))
    self.assertFalse(
        jnp.array_equal(saved_batch['inputs'],
                        saved_batch_diff_epoch['inputs']))
    self.assertFalse(
        jnp.array_equal(saved_batch['targets'],
                        saved_batch_diff_epoch['targets']))
    self.assertTrue(
        jnp.array_equal(saved_batch_eval['inputs'],
                        saved_batch_eval_same_epoch['inputs']))

    # Check shapes
    expected_shape = jnp.array([
        batch_size, hps.input_shape[0], hps.input_shape[1], hps.input_shape[2]
    ])
    expected_shape_eval = jnp.array([
        eval_batch_size, hps.input_shape[0],
        hps.input_shape[1], hps.input_shape[2],
    ])
    self.assertTrue(
        jnp.array_equal(saved_batch['inputs'].shape, expected_shape))
    self.assertTrue(
        jnp.array_equal(saved_batch_eval['inputs'].shape, expected_shape_eval))

    expected_target_shape = jnp.array(

# 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

    self.assertTrue(
        jnp.array_equal(saved_batch['inputs'],
                        saved_batch_same_epoch['inputs']))
    self.assertTrue(
        jnp.array_equal(saved_batch['targets'],
                        saved_batch_same_epoch['targets']))
    self.assertFalse(
        jnp.array_equal(saved_batch['inputs'],
                        saved_batch_diff_epoch['inputs']))
    self.assertFalse(
        jnp.array_equal(saved_batch['targets'],
                        saved_batch_diff_epoch['targets']))
    self.assertTrue(
        jnp.array_equal(saved_batch_eval['inputs'],
                        saved_batch_eval_same_epoch['inputs']))

    # Check shapes
    expected_shape = jnp.array([
        batch_size, hps.input_shape[0], hps.input_shape[1], hps.input_shape[2]
    ])
    expected_shape_eval = jnp.array([
        eval_batch_size, hps.input_shape[0],
        hps.input_shape[1], hps.input_shape[2],
    ])
    self.assertTrue(
        jnp.array_equal(saved_batch['inputs'].shape, expected_shape))
    self.assertTrue(
        jnp.array_equal(saved_batch_eval['inputs'].shape, expected_shape_eval))

    expected_target_shape = jnp.array(

        [batch_size, get_dataset_hparams(ds)['output_shape'][-1]])