# 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/pruning.py
# def weight_magnitude(weights):
# def prune(
#     model,
#     pruning_rate,
#     saliency_fn = weight_magnitude,
#     mask = None,
#     compare_fn = jnp.greater):


from typing import Mapping, Optional, Sequence
from absl.testing import absltest
from rigl.experimental.jax.pruning import masked
from rigl.experimental.jax.pruning import pruning
import flax
import jax
import jax.numpy as jnp
  """Tests the flax layer pruning module."""

  def setUp(self):
    super().setUp()
    self._rng = jax.random.PRNGKey(42)
    self._batch_size = 2
    self._input_shape = ((self._batch_size, 28, 28, 1), jnp.float32)
    self._input = jnp.ones(*self._input_shape)

    _, initial_params = MaskedDense.init_by_shape(self._rng,
                                                  (self._input_shape,))
    self._masked_model = flax.deprecated.nn.Model(MaskedDense, initial_params)

    _, initial_params = MaskedTwoLayerDense.init_by_shape(
        self._rng, (self._input_shape,))
    self._masked_model_twolayer = flax.deprecated.nn.Model(
        MaskedTwoLayerDense, initial_params)

    _, initial_params = MaskedConv.init_by_shape(self._rng,
                                                 (self._input_shape,))
    self._masked_conv_model = flax.deprecated.nn.Model(MaskedConv,
                                                       initial_params)

    _, initial_params = MaskedTwoLayerConv.init_by_shape(
        self._rng, (self._input_shape,))
    self._masked_conv_model_twolayer = flax.deprecated.nn.Model(
        MaskedTwoLayerConv, initial_params)

  def test_prune_single_layer_dense_no_mask(self):
    """Tests pruning of single dense layer without an existing mask."""

# 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/pruning.py
# def weight_magnitude(weights):
# def prune(
#     model,
#     pruning_rate,
#     saliency_fn = weight_magnitude,
#     mask = None,
#     compare_fn = jnp.greater):


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

  """Tests the flax layer pruning module."""

  def setUp(self):
    super().setUp()
    self._rng = jax.random.PRNGKey(42)
    self._batch_size = 2
    self._input_shape = ((self._batch_size, 28, 28, 1), jnp.float32)
    self._input = jnp.ones(*self._input_shape)

    _, initial_params = MaskedDense.init_by_shape(self._rng,
                                                  (self._input_shape,))
    self._masked_model = flax.deprecated.nn.Model(MaskedDense, initial_params)

    _, initial_params = MaskedTwoLayerDense.init_by_shape(
        self._rng, (self._input_shape,))
    self._masked_model_twolayer = flax.deprecated.nn.Model(
        MaskedTwoLayerDense, initial_params)

    _, initial_params = MaskedConv.init_by_shape(self._rng,
                                                 (self._input_shape,))
    self._masked_conv_model = flax.deprecated.nn.Model(MaskedConv,
                                                       initial_params)

    _, initial_params = MaskedTwoLayerConv.init_by_shape(
        self._rng, (self._input_shape,))
    self._masked_conv_model_twolayer = flax.deprecated.nn.Model(
        MaskedTwoLayerConv, initial_params)

  def test_prune_single_layer_dense_no_mask(self):
    """Tests pruning of single dense layer without an existing mask."""

    pruned_mask = pruning.prune(self._masked_model, 0.5)