# 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/utils/utils.py
# def cross_entropy_loss(log_softmax_logits,
#                        labels):
# def compute_metrics(logits,
#                     labels):
# def _np_converter(obj):
# def dump_dict_json(data_dict, path):
# def count_param(model,
#                 param_names):
# def cosine_similarity(a, b):
# def param_as_array(params):
# def cosine_similarity_model(initial_model,
#                             current_model):
# def vector_difference_norm_model(initial_model,
#                                  current_model):
# def pairwise_longest(iterable):
# T = TypeVar('T')


from typing import Sequence
from absl.testing import absltest
from absl.testing import parameterized
from rigl.experimental.jax.models import mnist_fc
from rigl.experimental.jax.utils import utils
import flax
import jax
import jax.numpy as jnp
    self._num_classes = 10
    self._batch_size = 2
    self._input_len = 28*28*1
    self._input_shape = ((self._batch_size, self._input_len), jnp.float32)
    self._input = jnp.zeros((self._batch_size, self._input_len), jnp.float32)
    self._param_count = 1e7

  def test_output_shapes(self):
    """Tests the output shape from the model."""
    with flax.deprecated.nn.stateful() as initial_state:
      _, initial_params = mnist_fc.MNISTFC.init_by_shape(
          self._rng, (self._input_shape,), num_classes=self._num_classes)
      model = flax.deprecated.nn.Model(mnist_fc.MNISTFC, 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_masks_depth(self):
    """Tests a model with an invalid mask."""
    invalid_masks = {
        'MaskedModule_0': {
            'kernel':
                jnp.zeros((self._batch_size, 5 * 5 * 16))
        }
    }

    with self.assertRaisesRegex(
        ValueError, 'Mask is invalid for model.'):

# 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/utils/utils.py
# def cross_entropy_loss(log_softmax_logits,
#                        labels):
# def compute_metrics(logits,
#                     labels):
# def _np_converter(obj):
# def dump_dict_json(data_dict, path):
# def count_param(model,
#                 param_names):
# def cosine_similarity(a, b):
# def param_as_array(params):
# def cosine_similarity_model(initial_model,
#                             current_model):
# def vector_difference_norm_model(initial_model,
#                                  current_model):
# def pairwise_longest(iterable):
# T = TypeVar('T')


# Path: rigl/experimental/jax/models/mnist_fc_test.py
from typing import Sequence
from absl.testing import absltest
from absl.testing import parameterized
from rigl.experimental.jax.models import mnist_fc
from rigl.experimental.jax.utils import utils
import flax
import jax
import jax.numpy as jnp

    self._num_classes = 10
    self._batch_size = 2
    self._input_len = 28*28*1
    self._input_shape = ((self._batch_size, self._input_len), jnp.float32)
    self._input = jnp.zeros((self._batch_size, self._input_len), jnp.float32)
    self._param_count = 1e7

  def test_output_shapes(self):
    """Tests the output shape from the model."""
    with flax.deprecated.nn.stateful() as initial_state:
      _, initial_params = mnist_fc.MNISTFC.init_by_shape(
          self._rng, (self._input_shape,), num_classes=self._num_classes)
      model = flax.deprecated.nn.Model(mnist_fc.MNISTFC, 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_masks_depth(self):
    """Tests a model with an invalid mask."""
    invalid_masks = {
        'MaskedModule_0': {
            'kernel':
                jnp.zeros((self._batch_size, 5 * 5 * 16))
        }
    }

    with self.assertRaisesRegex(
        ValueError, 'Mask is invalid for model.'):

      mnist_fc.MNISTFC.init_by_shape(