# Path: rigl/experimental/jax/datasets/cifar10.py
# class CIFAR10Dataset(dataset_base.ImageDataset):
#   NAME: str = 'cifar10'
#   MEAN_RGB: Sequence[float] = [0.4914 * 255, 0.4822 * 255, 0.4465 * 255]
#   STDDEV_RGB: Sequence[float] = [0.2470 * 255, 0.2435 * 255, 0.2616 * 255]
#   def __init__(self,
#                batch_size,
#                batch_size_test,
#                shuffle_buffer_size = 1024,
#                seed = 42):
#   def preprocess(
#       self, data):


from absl.testing import absltest
from rigl.experimental.jax.datasets import cifar10
import numpy as np
    with self.subTest(name='LabelShape'):
      self.assertLen(label, self._batch_size_test)

    with self.subTest(name='LabelType'):
      self.assertTrue(np.issubdtype(label.dtype, np.int))

    with self.subTest(name='LabelValues'):
      self.assertTrue((label >= 0).all() and
                      (label <= self._dataset.num_classes).all())

  def test_train_data_length(self):
    """Tests length of training dataset."""
    total_count = 0
    for batch in self._dataset.get_train():
      total_count += len(batch['label'])

    self.assertEqual(total_count, self._dataset.get_train_len())

  def test_test_data_length(self):
    """Tests length of test dataset."""
    total_count = 0
    for batch in self._dataset.get_test():
      total_count += len(batch['label'])

    self.assertEqual(total_count, self._dataset.get_test_len())

  def test_dataset_nonevenly_divisible_batch_size(self):
    """Tests non-evenly divisible test batch size."""
    with self.assertRaisesRegex(
        ValueError, 'Test data not evenly divisible by batch size: .*'):

# Path: rigl/experimental/jax/datasets/cifar10.py
# class CIFAR10Dataset(dataset_base.ImageDataset):
#   NAME: str = 'cifar10'
#   MEAN_RGB: Sequence[float] = [0.4914 * 255, 0.4822 * 255, 0.4465 * 255]
#   STDDEV_RGB: Sequence[float] = [0.2470 * 255, 0.2435 * 255, 0.2616 * 255]
#   def __init__(self,
#                batch_size,
#                batch_size_test,
#                shuffle_buffer_size = 1024,
#                seed = 42):
#   def preprocess(
#       self, data):


# Path: rigl/experimental/jax/datasets/cifar10_test.py
from absl.testing import absltest
from rigl.experimental.jax.datasets import cifar10
import numpy as np

    with self.subTest(name='LabelShape'):
      self.assertLen(label, self._batch_size_test)

    with self.subTest(name='LabelType'):
      self.assertTrue(np.issubdtype(label.dtype, np.int))

    with self.subTest(name='LabelValues'):
      self.assertTrue((label >= 0).all() and
                      (label <= self._dataset.num_classes).all())

  def test_train_data_length(self):
    """Tests length of training dataset."""
    total_count = 0
    for batch in self._dataset.get_train():
      total_count += len(batch['label'])

    self.assertEqual(total_count, self._dataset.get_train_len())

  def test_test_data_length(self):
    """Tests length of test dataset."""
    total_count = 0
    for batch in self._dataset.get_test():
      total_count += len(batch['label'])

    self.assertEqual(total_count, self._dataset.get_test_len())

  def test_dataset_nonevenly_divisible_batch_size(self):
    """Tests non-evenly divisible test batch size."""
    with self.assertRaisesRegex(
        ValueError, 'Test data not evenly divisible by batch size: .*'):

      self._dataset = cifar10.CIFAR10Dataset(