# Path: rigl/experimental/jax/pruning/mask_factory.py
# MASK_TYPES: Mapping[str, MaskFnType] = {
#     'random':
#         masked.shuffled_mask,
#     'per_neuron':
#         masked.shuffled_neuron_mask,
#     'per_neuron_no_input_ablation':
#         masked.shuffled_neuron_no_input_ablation_mask,
#     'symmetric':
#         masked.symmetric_mask,
# }
# def create_mask(mask_type, base_model,
#                 rng, sparsity,
#                 **kwargs):
# 
# Path: rigl/experimental/jax/pruning/masked.py
# def masked(module, mask):
#   """Convenience function for masking a FLAX module with MaskedModule."""
#   return MaskedModule.partial(wrapped_module=module, mask=mask)


from typing import Mapping, Optional
from absl.testing import absltest
from absl.testing import parameterized
from rigl.experimental.jax.pruning import mask_factory
from rigl.experimental.jax.pruning import masked
import flax
import jax
import jax.numpy as jnp
            inputs,
            mask = None):
    inputs = inputs.reshape(inputs.shape[0], -1)

    return masked.MaskedModule(
        inputs,
        features=self.NUM_FEATURES,
        wrapped_module=flax.deprecated.nn.Dense,
        mask=mask['MaskedModule_0'] if mask else None)


class MaskFactoryTest(parameterized.TestCase):

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

    _, initial_params = MaskedDense.init_by_shape(self._rng,
                                                  (self._input_shape,))
    # Use the same initialization for both masked/unmasked models.
    self._model = flax.deprecated.nn.Model(MaskedDense, initial_params)

  def _create_mask(self, mask_type):
    return mask_factory.create_mask(
        mask_type, self._model,
        self._rng, self._sparsity)


# Path: rigl/experimental/jax/pruning/mask_factory.py
# MASK_TYPES: Mapping[str, MaskFnType] = {
#     'random':
#         masked.shuffled_mask,
#     'per_neuron':
#         masked.shuffled_neuron_mask,
#     'per_neuron_no_input_ablation':
#         masked.shuffled_neuron_no_input_ablation_mask,
#     'symmetric':
#         masked.symmetric_mask,
# }
# def create_mask(mask_type, base_model,
#                 rng, sparsity,
#                 **kwargs):
# 
# Path: rigl/experimental/jax/pruning/masked.py
# def masked(module, mask):
#   """Convenience function for masking a FLAX module with MaskedModule."""
#   return MaskedModule.partial(wrapped_module=module, mask=mask)


# Path: rigl/experimental/jax/pruning/mask_factory_test.py
from typing import Mapping, Optional
from absl.testing import absltest
from absl.testing import parameterized
from rigl.experimental.jax.pruning import mask_factory
from rigl.experimental.jax.pruning import masked
import flax
import jax
import jax.numpy as jnp

            inputs,
            mask = None):
    inputs = inputs.reshape(inputs.shape[0], -1)

    return masked.MaskedModule(
        inputs,
        features=self.NUM_FEATURES,
        wrapped_module=flax.deprecated.nn.Dense,
        mask=mask['MaskedModule_0'] if mask else None)


class MaskFactoryTest(parameterized.TestCase):

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

    _, initial_params = MaskedDense.init_by_shape(self._rng,
                                                  (self._input_shape,))
    # Use the same initialization for both masked/unmasked models.
    self._model = flax.deprecated.nn.Model(MaskedDense, initial_params)

  def _create_mask(self, mask_type):
    return mask_factory.create_mask(
        mask_type, self._model,
        self._rng, self._sparsity)


  @parameterized.parameters(*mask_factory.MASK_TYPES.keys())