# 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/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_fc.py
# def feature_dim_for_param(input_len,
#                           param_count,
#                           depth,
#                           depth_mult = 2.):
#   def apply(self,
#             inputs,
#             num_classes,
#             features = (32, 32),
#             train=True,
#             init_fn = flax.deprecated.nn.initializers.kaiming_normal,
#             activation_fn = flax.deprecated.nn.relu,
#             masks = None,
#             masked_layer_indices = None,
#             dropout_rate = 0.):
# class MNISTFC(flax.deprecated.nn.Module):


from typing import Any, Callable, Mapping, Sequence, Tuple, Type
from rigl.experimental.jax.models import cifar10_cnn
from rigl.experimental.jax.models import mnist_cnn
from rigl.experimental.jax.models import mnist_fc
import flax
import jax.numpy as jnp
# coding=utf-8
# Copyright 2022 RigL Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# 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
"""Factory for neural network models.

Attributes:
  MODELS: A list of the models that can be created.
"""



MODELS: Mapping[str, Type[flax.deprecated.nn.Model]] = {
    'MNIST_CNN': mnist_cnn.MNISTCNN,
    'MNIST_FC': mnist_fc.MNISTFC,

# 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/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_fc.py
# def feature_dim_for_param(input_len,
#                           param_count,
#                           depth,
#                           depth_mult = 2.):
#   def apply(self,
#             inputs,
#             num_classes,
#             features = (32, 32),
#             train=True,
#             init_fn = flax.deprecated.nn.initializers.kaiming_normal,
#             activation_fn = flax.deprecated.nn.relu,
#             masks = None,
#             masked_layer_indices = None,
#             dropout_rate = 0.):
# class MNISTFC(flax.deprecated.nn.Module):


# Path: rigl/experimental/jax/models/model_factory.py
from typing import Any, Callable, Mapping, Sequence, Tuple, Type
from rigl.experimental.jax.models import cifar10_cnn
from rigl.experimental.jax.models import mnist_cnn
from rigl.experimental.jax.models import mnist_fc
import flax
import jax.numpy as jnp

# coding=utf-8
# Copyright 2022 RigL Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# 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
"""Factory for neural network models.

Attributes:
  MODELS: A list of the models that can be created.
"""



MODELS: Mapping[str, Type[flax.deprecated.nn.Model]] = {
    'MNIST_CNN': mnist_cnn.MNISTCNN,
    'MNIST_FC': mnist_fc.MNISTFC,

    'CIFAR10_CNN': cifar10_cnn.CIFAR10CNN,