# 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/utils/utils.py
# def cross_entropy_loss(log_softmax_logits,
#                        labels):
# def compute_metrics(logits,
#                     labels):
# def _np_converter(obj):
# def dump_dict_json(data_dict, path):
# def count_param(model,
#                 param_names):
# def cosine_similarity(a, b):
# def param_as_array(params):
# def cosine_similarity_model(initial_model,
#                             current_model):
# def vector_difference_norm_model(initial_model,
#                                  current_model):
# def pairwise_longest(iterable):
# T = TypeVar('T')


import functools
import math
import operator
import jax.numpy as jnp
import numpy as np
from typing import Dict, Optional, Union
from rigl.experimental.jax.pruning import masked
from rigl.experimental.jax.utils import utils
  mask_stats['zeroed_neurons'] = int(zeroed_count)

  mask_stats['permutations'] = functools.reduce(
      operator.mul, (np.math.factorial(t) for t in unique_counts))
  mask_stats['unique_neurons'] = len(unique_counts)

  return mask_stats


def count_permutations_mask(mask):
  """Calculates the number of permutations for a given model mask.

  Args:
    mask: Model masks to check, similar to Model.params.

  Returns:
   A dictionary with stats on the permutation structure of a mask, including
   the number of symmetric permutations of the mask, number of unique mask
   columns, count of the zeroed out (structurally pruned) neurons, and total
   number of neurons/filters.
  """
  sum_keys = ('total_neurons', 'unique_neurons', 'zeroed_neurons')
  product_keys = ('permutations',)

  # Count permutation stats for each pairwise set of layers.
  # Note: I tried doing this with more_itertools.pairwise/itertools.chain, but
  # there is a type conflict in passing iterators of different types to
  # itertools.chain.
  counts = [
      count_permutations_mask_layer(layer, next_layer)

# 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/utils/utils.py
# def cross_entropy_loss(log_softmax_logits,
#                        labels):
# def compute_metrics(logits,
#                     labels):
# def _np_converter(obj):
# def dump_dict_json(data_dict, path):
# def count_param(model,
#                 param_names):
# def cosine_similarity(a, b):
# def param_as_array(params):
# def cosine_similarity_model(initial_model,
#                             current_model):
# def vector_difference_norm_model(initial_model,
#                                  current_model):
# def pairwise_longest(iterable):
# T = TypeVar('T')


# Path: rigl/experimental/jax/pruning/symmetry.py
import functools
import math
import operator
import jax.numpy as jnp
import numpy as np
from typing import Dict, Optional, Union
from rigl.experimental.jax.pruning import masked
from rigl.experimental.jax.utils import utils

  mask_stats['zeroed_neurons'] = int(zeroed_count)

  mask_stats['permutations'] = functools.reduce(
      operator.mul, (np.math.factorial(t) for t in unique_counts))
  mask_stats['unique_neurons'] = len(unique_counts)

  return mask_stats


def count_permutations_mask(mask):
  """Calculates the number of permutations for a given model mask.

  Args:
    mask: Model masks to check, similar to Model.params.

  Returns:
   A dictionary with stats on the permutation structure of a mask, including
   the number of symmetric permutations of the mask, number of unique mask
   columns, count of the zeroed out (structurally pruned) neurons, and total
   number of neurons/filters.
  """
  sum_keys = ('total_neurons', 'unique_neurons', 'zeroed_neurons')
  product_keys = ('permutations',)

  # Count permutation stats for each pairwise set of layers.
  # Note: I tried doing this with more_itertools.pairwise/itertools.chain, but
  # there is a type conflict in passing iterators of different types to
  # itertools.chain.
  counts = [
      count_permutations_mask_layer(layer, next_layer)

      for layer, next_layer in utils.pairwise_longest(mask.values())