# Path: rigl/experimental/jax/models/model_factory.py
# MODELS: Mapping[str, Type[flax.deprecated.nn.Model]] = {
#     'MNIST_CNN': mnist_cnn.MNISTCNN,
#     'MNIST_FC': mnist_fc.MNISTFC,
#     'CIFAR10_CNN': cifar10_cnn.CIFAR10CNN,
# }
# def create_model(
#     name, rng,
#     input_specs, **kwargs
# ):
# def update_model(model,
#                  **kwargs):


from absl.testing import absltest
from absl.testing import parameterized
from rigl.experimental.jax.models import model_factory
import flax
import jax
import jax.numpy as jnp
# 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.model_factory."""



class ModelCommonTest(parameterized.TestCase):
  """Tests the model factory."""

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

  def _create_model(self, model_name):
    return model_factory.create_model(
        model_name,
        self._rng, (self._input_shape,),
        num_classes=self._num_classes)


# Path: rigl/experimental/jax/models/model_factory.py
# MODELS: Mapping[str, Type[flax.deprecated.nn.Model]] = {
#     'MNIST_CNN': mnist_cnn.MNISTCNN,
#     'MNIST_FC': mnist_fc.MNISTFC,
#     'CIFAR10_CNN': cifar10_cnn.CIFAR10CNN,
# }
# def create_model(
#     name, rng,
#     input_specs, **kwargs
# ):
# def update_model(model,
#                  **kwargs):


# Path: rigl/experimental/jax/models/model_factory_test.py
from absl.testing import absltest
from absl.testing import parameterized
from rigl.experimental.jax.models import model_factory
import flax
import jax
import jax.numpy as jnp

# 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.model_factory."""



class ModelCommonTest(parameterized.TestCase):
  """Tests the model factory."""

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

  def _create_model(self, model_name):
    return model_factory.create_model(
        model_name,
        self._rng, (self._input_shape,),
        num_classes=self._num_classes)


  @parameterized.parameters(*model_factory.MODELS.keys())