# Path: rigl/experimental/jax/datasets/dataset_base.py
# class Dataset(metaclass=abc.ABCMeta):
# class ImageDataset(Dataset):
#   DATAKEY: Optional[str] = None
#   LABELKEY: str = 'label'
#   DATAKEY = 'image'
#   def __init__(self,
#                name,
#                batch_size,
#                batch_size_test,
#                shuffle_buffer_size,
#                prefetch_size = 1,
#                seed = None):  # pytype: disable=annotation-type-mismatch
#   def _dataset_dir(self):
#   def get_train(self):
#   def get_train_len(self):
#   def get_test(self):
#   def get_test_len(self):
#   def preprocess(
#       self, data):
#   def augment(
#       self, data):
#   def preprocess(
#       self, data):
# 
# Path: rigl/experimental/jax/datasets/dataset_factory.py
# DATASETS: Mapping[str, Type[dataset_base.Dataset]] = {
#     'MNIST': mnist.MNISTDataset,
#     'CIFAR10': cifar10.CIFAR10Dataset,
# }
# def create_dataset(name, *args, **kwargs):


from absl.testing import absltest
from absl.testing import parameterized
from rigl.experimental.jax.datasets import dataset_base
from rigl.experimental.jax.datasets import dataset_factory
import numpy as np
# See the License for the specific language governing permissions and
# limitations under the License.

# Lint as: python3
"""Tests for weight_symmetry.datasets.dataset_common."""


class DatasetCommonTest(parameterized.TestCase):

  def setUp(self):
    super().setUp()
    self._batch_size = 32
    self._batch_size_test = 10
    self._shuffle_buffer_size = 128

  def _create_dataset(self, dataset_name):
    """Helper function for creating a dataset."""
    return dataset_factory.create_dataset(
        dataset_name,
        self._batch_size,
        self._batch_size_test,
        shuffle_buffer_size=self._shuffle_buffer_size)

  def test_dataset_supported(self):
    """Tests supported datasets."""
    for dataset_name in dataset_factory.DATASETS:
      dataset = self._create_dataset(dataset_name)

      self.assertIsInstance(dataset, dataset_base.Dataset)


# Path: rigl/experimental/jax/datasets/dataset_base.py
# class Dataset(metaclass=abc.ABCMeta):
# class ImageDataset(Dataset):
#   DATAKEY: Optional[str] = None
#   LABELKEY: str = 'label'
#   DATAKEY = 'image'
#   def __init__(self,
#                name,
#                batch_size,
#                batch_size_test,
#                shuffle_buffer_size,
#                prefetch_size = 1,
#                seed = None):  # pytype: disable=annotation-type-mismatch
#   def _dataset_dir(self):
#   def get_train(self):
#   def get_train_len(self):
#   def get_test(self):
#   def get_test_len(self):
#   def preprocess(
#       self, data):
#   def augment(
#       self, data):
#   def preprocess(
#       self, data):
# 
# Path: rigl/experimental/jax/datasets/dataset_factory.py
# DATASETS: Mapping[str, Type[dataset_base.Dataset]] = {
#     'MNIST': mnist.MNISTDataset,
#     'CIFAR10': cifar10.CIFAR10Dataset,
# }
# def create_dataset(name, *args, **kwargs):


# Path: rigl/experimental/jax/datasets/dataset_factory_test.py
from absl.testing import absltest
from absl.testing import parameterized
from rigl.experimental.jax.datasets import dataset_base
from rigl.experimental.jax.datasets import dataset_factory
import numpy as np

# See the License for the specific language governing permissions and
# limitations under the License.

# Lint as: python3
"""Tests for weight_symmetry.datasets.dataset_common."""


class DatasetCommonTest(parameterized.TestCase):

  def setUp(self):
    super().setUp()
    self._batch_size = 32
    self._batch_size_test = 10
    self._shuffle_buffer_size = 128

  def _create_dataset(self, dataset_name):
    """Helper function for creating a dataset."""
    return dataset_factory.create_dataset(
        dataset_name,
        self._batch_size,
        self._batch_size_test,
        shuffle_buffer_size=self._shuffle_buffer_size)

  def test_dataset_supported(self):
    """Tests supported datasets."""
    for dataset_name in dataset_factory.DATASETS:
      dataset = self._create_dataset(dataset_name)

      self.assertIsInstance(dataset, dataset_base.Dataset)


  @parameterized.parameters(*dataset_factory.DATASETS.keys())