# 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
    mask_layer_1_sparsity = masked.mask_sparsity(pruned_mask['MaskedModule_1'])

    with self.subTest(name='test_mask_layer1_param_not_none'):
      self.assertNotEmpty(pruned_mask['MaskedModule_0']['kernel'])

    with self.subTest(name='test_mask_layer2_param_not_none'):
      self.assertNotEmpty(pruned_mask['MaskedModule_1']['kernel'])

    with self.subTest(name='test_mask_layer_0_sparsity'):
      self.assertEqual(mask_layer_0_sparsity, 0.)

    with self.subTest(name='test_mask_layer_1_sparsity'):
      self.assertAlmostEqual(mask_layer_1_sparsity, 0.5, places=3)

  def test_prune_one_layer_conv_no_mask(self):
    """Tests pruning of model with one conv. layer without an existing mask."""
    pruned_mask = pruning.prune(self._masked_conv_model, 0.5)
    mask_sparsity = masked.mask_sparsity(pruned_mask)

    with self.subTest(name='test_mask_param_not_none'):
      self.assertNotEmpty(pruned_mask['MaskedModule_0']['kernel'])

    with self.subTest(name='test_mask_sparsity'):
      self.assertAlmostEqual(mask_sparsity, 0.5, places=1)

  def test_prune_one_layer_conv_with_mask(self):
    """Tests pruning of model with one conv. layer with an existing mask."""
    pruned_mask = pruning.prune(
        self._masked_conv_model,
        0.5,

# 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

    mask_layer_1_sparsity = masked.mask_sparsity(pruned_mask['MaskedModule_1'])

    with self.subTest(name='test_mask_layer1_param_not_none'):
      self.assertNotEmpty(pruned_mask['MaskedModule_0']['kernel'])

    with self.subTest(name='test_mask_layer2_param_not_none'):
      self.assertNotEmpty(pruned_mask['MaskedModule_1']['kernel'])

    with self.subTest(name='test_mask_layer_0_sparsity'):
      self.assertEqual(mask_layer_0_sparsity, 0.)

    with self.subTest(name='test_mask_layer_1_sparsity'):
      self.assertAlmostEqual(mask_layer_1_sparsity, 0.5, places=3)

  def test_prune_one_layer_conv_no_mask(self):
    """Tests pruning of model with one conv. layer without an existing mask."""
    pruned_mask = pruning.prune(self._masked_conv_model, 0.5)
    mask_sparsity = masked.mask_sparsity(pruned_mask)

    with self.subTest(name='test_mask_param_not_none'):
      self.assertNotEmpty(pruned_mask['MaskedModule_0']['kernel'])

    with self.subTest(name='test_mask_sparsity'):
      self.assertAlmostEqual(mask_sparsity, 0.5, places=1)

  def test_prune_one_layer_conv_with_mask(self):
    """Tests pruning of model with one conv. layer with an existing mask."""
    pruned_mask = pruning.prune(
        self._masked_conv_model,
        0.5,

        mask=masked.shuffled_mask(self._masked_model, self._rng, 0.95))