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


from absl.testing import absltest
from rigl.experimental.jax.models import cifar10_cnn
import flax
import jax
import jax.numpy as jnp
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     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.models.cifar10_cnn."""



class CIFAR10CNNTest(absltest.TestCase):
  """Tests the CIFAR10CNN 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, 32, 32, 3), 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:

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


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

# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     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.models.cifar10_cnn."""



class CIFAR10CNNTest(absltest.TestCase):
  """Tests the CIFAR10CNN 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, 32, 32, 3), 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 = cifar10_cnn.CIFAR10CNN.init_by_shape(