# 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)


import collections
import flax
import jax.numpy as jnp
from typing import Any, Callable, Mapping, Optional, Union
from rigl.experimental.jax.pruning import masked


def prune(
    model,
    pruning_rate,
    saliency_fn = weight_magnitude,
    mask = None,
    compare_fn = jnp.greater):
  """Returns a mask for a model where the params in each layer are pruned using a saliency function.

  Args:
    model: The model to create a pruning mask for.
    pruning_rate: The fraction of lowest magnitude saliency weights that are
      pruned. If a float, the same rate is used for all layers, otherwise if it
      is a mapping, it must contain a rate for all masked layers in the model.
    saliency_fn: A function that returns a float number used to rank
      the importance of individual weights in the layer.
    mask: If the model has an existing mask, the mask will be applied before
      pruning the model.
    compare_fn: A pairwise operator to compare saliency with threshold, and
      return True if the saliency indicates the value should not be masked.

  Returns:
    A pruned mask for the given model.
  """
  if not mask:
    mask = masked.simple_mask(model, jnp.ones, masked.WEIGHT_PARAM_NAMES)

  if not isinstance(pruning_rate, collections.Mapping):
    pruning_rate_dict = {}

# 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
import collections
import flax
import jax.numpy as jnp
from typing import Any, Callable, Mapping, Optional, Union
from rigl.experimental.jax.pruning import masked



def prune(
    model,
    pruning_rate,
    saliency_fn = weight_magnitude,
    mask = None,
    compare_fn = jnp.greater):
  """Returns a mask for a model where the params in each layer are pruned using a saliency function.

  Args:
    model: The model to create a pruning mask for.
    pruning_rate: The fraction of lowest magnitude saliency weights that are
      pruned. If a float, the same rate is used for all layers, otherwise if it
      is a mapping, it must contain a rate for all masked layers in the model.
    saliency_fn: A function that returns a float number used to rank
      the importance of individual weights in the layer.
    mask: If the model has an existing mask, the mask will be applied before
      pruning the model.
    compare_fn: A pairwise operator to compare saliency with threshold, and
      return True if the saliency indicates the value should not be masked.

  Returns:
    A pruned mask for the given model.
  """
  if not mask:
    mask = masked.simple_mask(model, jnp.ones, masked.WEIGHT_PARAM_NAMES)

  if not isinstance(pruning_rate, collections.Mapping):
    pruning_rate_dict = {}

    for param_name, _ in masked.iterate_mask(mask):