# 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
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Lint as: python3
"""Tests for weight_symmetry.pruning.masked."""




class Dense(flax.deprecated.nn.Module):
  """Single-layer Dense Non-Masked Network."""

  NUM_FEATURES: int = 32

  def apply(self, inputs):
    inputs = inputs.reshape(inputs.shape[0], -1)
    return flax.deprecated.nn.Dense(inputs, features=self.NUM_FEATURES)


class MaskedDense(flax.deprecated.nn.Module):
  """Single-layer Dense Masked Network."""

  NUM_FEATURES: int = 32

  def apply(self,
            inputs,
            mask = None):
    inputs = inputs.reshape(inputs.shape[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

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Lint as: python3
"""Tests for weight_symmetry.pruning.masked."""




class Dense(flax.deprecated.nn.Module):
  """Single-layer Dense Non-Masked Network."""

  NUM_FEATURES: int = 32

  def apply(self, inputs):
    inputs = inputs.reshape(inputs.shape[0], -1)
    return flax.deprecated.nn.Dense(inputs, features=self.NUM_FEATURES)


class MaskedDense(flax.deprecated.nn.Module):
  """Single-layer Dense Masked Network."""

  NUM_FEATURES: int = 32

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


    return masked.MaskedModule(