# 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, Sequence
from absl.testing import absltest
from absl.testing import parameterized
from rigl.experimental.jax.pruning import masked
import flax
import jax
import jax.numpy as jnp
import numpy as np

    with self.subTest(name='shuffled_empty_mask_layer2'):
      self.assertIn('MaskedModule_1', mask)

    with self.subTest(name='shuffled_empty_mask_values_layer2'):
      self.assertTrue((mask['MaskedModule_1']['kernel'] == 1).all())

    masked_output = self._masked_model_twolayer(self._input, mask=mask)

    with self.subTest(name='shuffled_empty_dense_values'):
      self.assertTrue(
          jnp.isclose(masked_output, self._unmasked_output_twolayer).all())

    with self.subTest(name='shuffled_empty_mask_dense_shape'):
      self.assertSequenceEqual(masked_output.shape,
                               self._unmasked_output_twolayer.shape)

  def test_random_invalid_model(self):
    """Tests random mask with model containing no masked layers."""
    with self.assertRaisesRegex(
        ValueError, 'Model does not support masking, i.e. no layers are '
        'wrapped by a MaskedModule.'):
      masked.random_mask(self._unmasked_model, self._rng, 0.5)

  def test_random_invalid_sparsity(self):
    """Tests random mask with invalid sparsity."""

    with self.subTest(name='random_sparsity_too_small'):
      with self.assertRaisesRegex(
          ValueError, r'Given sparsity, -0.5, is not in range \[0, 1\]'):

# 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/masked_test.py
from typing import Mapping, Optional, Sequence
from absl.testing import absltest
from absl.testing import parameterized
from rigl.experimental.jax.pruning import masked
import flax
import jax
import jax.numpy as jnp
import numpy as np


    with self.subTest(name='shuffled_empty_mask_layer2'):
      self.assertIn('MaskedModule_1', mask)

    with self.subTest(name='shuffled_empty_mask_values_layer2'):
      self.assertTrue((mask['MaskedModule_1']['kernel'] == 1).all())

    masked_output = self._masked_model_twolayer(self._input, mask=mask)

    with self.subTest(name='shuffled_empty_dense_values'):
      self.assertTrue(
          jnp.isclose(masked_output, self._unmasked_output_twolayer).all())

    with self.subTest(name='shuffled_empty_mask_dense_shape'):
      self.assertSequenceEqual(masked_output.shape,
                               self._unmasked_output_twolayer.shape)

  def test_random_invalid_model(self):
    """Tests random mask with model containing no masked layers."""
    with self.assertRaisesRegex(
        ValueError, 'Model does not support masking, i.e. no layers are '
        'wrapped by a MaskedModule.'):
      masked.random_mask(self._unmasked_model, self._rng, 0.5)

  def test_random_invalid_sparsity(self):
    """Tests random mask with invalid sparsity."""

    with self.subTest(name='random_sparsity_too_small'):
      with self.assertRaisesRegex(
          ValueError, r'Given sparsity, -0.5, is not in range \[0, 1\]'):

        masked.random_mask(self._masked_model, self._rng, -0.5)