# Path: rigl/experimental/jax/datasets/mnist.py
# class MNISTDataset(dataset_base.ImageDataset):
#   NAME: str = 'mnist'
#   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 mnist
import numpy as np
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# 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.

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



class MNISTDatasetTest(absltest.TestCase):
  """Test cases for MNIST Dataset."""

  def setUp(self):
    """Common setup routines/variables for test cases."""
    super().setUp()
    self._batch_size = 16
    self._batch_size_test = 10
    self._shuffle_buffer_size = 8

    self._dataset = mnist.MNISTDataset(
        self._batch_size,
        batch_size_test=self._batch_size_test,
        shuffle_buffer_size=self._shuffle_buffer_size)

  def test_create_dataset(self):
    """Tests creation of dataset."""

# Path: rigl/experimental/jax/datasets/mnist.py
# class MNISTDataset(dataset_base.ImageDataset):
#   NAME: str = 'mnist'
#   def __init__(self,
#                batch_size,
#                batch_size_test,
#                shuffle_buffer_size = 1024,
#                seed = 42):
#   def preprocess(
#       self, data):


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

#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# 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.

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



class MNISTDatasetTest(absltest.TestCase):
  """Test cases for MNIST Dataset."""

  def setUp(self):
    """Common setup routines/variables for test cases."""
    super().setUp()
    self._batch_size = 16
    self._batch_size_test = 10
    self._shuffle_buffer_size = 8

    self._dataset = mnist.MNISTDataset(
        self._batch_size,
        batch_size_test=self._batch_size_test,
        shuffle_buffer_size=self._shuffle_buffer_size)

  def test_create_dataset(self):
    """Tests creation of dataset."""

    self.assertIsInstance(self._dataset, mnist.MNISTDataset)