# Path: rigl/experimental/jax/models/mnist_cnn.py
# class MNISTCNN(flax.deprecated.nn.Module):
#   def apply(self,
#             inputs,
#             num_classes,
#             filter_shape = (5, 5),
#             filters = (16, 32),
#             dense_size = 64,
#             train=True,
#             init_fn = flax.deprecated.nn.initializers.kaiming_normal,
#             activation_fn = flax.deprecated.nn.relu,
#             masks = None,
#             masked_layer_indices = None):


from absl.testing import absltest
from rigl.experimental.jax.models import mnist_cnn
import flax
import jax
import jax.numpy as jnp
"""Tests for weight_symmetry.models.mnist_cnn."""



class MNISTCNNTest(absltest.TestCase):
  """Tests the MNISTCNN model."""

  def setUp(self):
    super().setUp()
    self._rng = jax.random.PRNGKey(42)
    self._num_classes = 10
    self._batch_size = 2
    self._input_shape = ((self._batch_size, 28, 28, 1), jnp.float32)
    self._input = jnp.zeros(*self._input_shape)

  def test_output_shapes(self):
    """Tests the output shapes of the model."""
    with flax.deprecated.nn.stateful() as initial_state:
      _, initial_params = mnist_cnn.MNISTCNN.init_by_shape(
          self._rng, (self._input_shape,), num_classes=self._num_classes)
      model = flax.deprecated.nn.Model(mnist_cnn.MNISTCNN, initial_params)

    with flax.deprecated.nn.stateful(initial_state, mutable=False):
      logits = model(self._input, num_classes=self._num_classes, train=False)

    self.assertTupleEqual(logits.shape, (self._batch_size, self._num_classes))

  def test_invalid_depth(self):
    """Tests model mask with the incorrect depth for the given model."""
    with self.assertRaisesRegex(ValueError, 'Input spatial size, '):

# Path: rigl/experimental/jax/models/mnist_cnn.py
# class MNISTCNN(flax.deprecated.nn.Module):
#   def apply(self,
#             inputs,
#             num_classes,
#             filter_shape = (5, 5),
#             filters = (16, 32),
#             dense_size = 64,
#             train=True,
#             init_fn = flax.deprecated.nn.initializers.kaiming_normal,
#             activation_fn = flax.deprecated.nn.relu,
#             masks = None,
#             masked_layer_indices = None):


# Path: rigl/experimental/jax/models/mnist_cnn_test.py
from absl.testing import absltest
from rigl.experimental.jax.models import mnist_cnn
import flax
import jax
import jax.numpy as jnp

"""Tests for weight_symmetry.models.mnist_cnn."""



class MNISTCNNTest(absltest.TestCase):
  """Tests the MNISTCNN model."""

  def setUp(self):
    super().setUp()
    self._rng = jax.random.PRNGKey(42)
    self._num_classes = 10
    self._batch_size = 2
    self._input_shape = ((self._batch_size, 28, 28, 1), jnp.float32)
    self._input = jnp.zeros(*self._input_shape)

  def test_output_shapes(self):
    """Tests the output shapes of the model."""
    with flax.deprecated.nn.stateful() as initial_state:
      _, initial_params = mnist_cnn.MNISTCNN.init_by_shape(
          self._rng, (self._input_shape,), num_classes=self._num_classes)
      model = flax.deprecated.nn.Model(mnist_cnn.MNISTCNN, initial_params)

    with flax.deprecated.nn.stateful(initial_state, mutable=False):
      logits = model(self._input, num_classes=self._num_classes, train=False)

    self.assertTupleEqual(logits.shape, (self._batch_size, self._num_classes))

  def test_invalid_depth(self):
    """Tests model mask with the incorrect depth for the given model."""
    with self.assertRaisesRegex(ValueError, 'Input spatial size, '):

      mnist_cnn.MNISTCNN.init_by_shape(