# Path: rigl/sparse_utils.py
# DEFAULT_ERK_SCALE = 1.0
# def mask_extract_name_fn(mask_name):
# def get_n_zeros(size, sparsity):
# def calculate_sparsity(masks):
# def get_mask_random_numpy(mask_shape, sparsity, random_state=None):
# def get_mask_random(mask, sparsity, dtype, random_state=None):
# def get_sparsities_erdos_renyi(all_masks,
#                                default_sparsity,
#                                custom_sparsity_map,
#                                include_kernel,
#                                extract_name_fn=mask_extract_name_fn,
#                                erk_power_scale=DEFAULT_ERK_SCALE):
# def get_sparsities_uniform(all_masks,
#                            default_sparsity,
#                            custom_sparsity_map,
#                            extract_name_fn=mask_extract_name_fn):
# def get_sparsities_str(all_masks, default_sparsity):
# def get_sparsities(all_masks,
#                    method,
#                    default_sparsity,
#                    custom_sparsity_map,
#                    extract_name_fn=mask_extract_name_fn,
#                    erk_power_scale=DEFAULT_ERK_SCALE):
# def get_mask_init_fn(all_masks,
#                      method,
#                      default_sparsity,
#                      custom_sparsity_map,
#                      mask_fn=get_mask_random,
#                      erk_power_scale=DEFAULT_ERK_SCALE,
#                      extract_name_fn=mask_extract_name_fn):
# def _get_kernel(layer):
# def get_stats(masked_layers,
#               default_sparsity=0.8,
#               method='erdos_renyi',
#               custom_sparsities=None,
#               is_debug=False,
#               width=1.,
#               first_layer_name='conv1',
#               last_layer_name='conv_preds',
#               param_size=32,
#               erk_power_scale=DEFAULT_ERK_SCALE):


import absl.testing.parameterized as parameterized
import numpy as np
import tensorflow.compat.v1 as tf
from rigl import sparse_utils
                 tf.get_variable(shape=(1, 1, 3), name='var3/mask')]
    custom_sparsity = {'var3': 0.8}
    sparsities = sparse_utils.get_sparsities(
        all_masks, 'erdos_renyi', default_sparsity, custom_sparsity)
    self.assertEqual(sparsities[all_masks[2].name], 0.8)

  @parameterized.parameters(((2, 3), (2, 3), 0.5),
                            ((1, 1, 2, 3), (1, 1, 2, 3), 0.3),
                            ((8, 6), (4, 3), 0.7),
                            ((80, 4), (20, 20), 0.8),
                            ((2, 6), (2, 3), 0.8))
  def testSparsityDictErdosRenyiSparsitiesScale(
      self, shape1, shape2, default_sparsity):
    _ = self._setup_session()
    all_masks = [tf.get_variable(shape=shape1, name='var1/mask'),
                 tf.get_variable(shape=shape2, name='var2/mask')]
    custom_sparsity = {}
    sparsities = sparse_utils.get_sparsities(
        all_masks, 'erdos_renyi', default_sparsity, custom_sparsity)
    sparsity1 = sparsities[all_masks[0].name]
    size1 = np.prod(shape1)
    sparsity2 = sparsities[all_masks[1].name]
    size2 = np.prod(shape2)
    # Ensure that total number of connections are similar.
    expected_zeros_uniform = (
        sparse_utils.get_n_zeros(size1, default_sparsity) +
        sparse_utils.get_n_zeros(size2, default_sparsity))
    # Ensure that total number of connections are similar.
    expected_zeros_current = (
        sparse_utils.get_n_zeros(size1, sparsity1) +

# Path: rigl/sparse_utils.py
# DEFAULT_ERK_SCALE = 1.0
# def mask_extract_name_fn(mask_name):
# def get_n_zeros(size, sparsity):
# def calculate_sparsity(masks):
# def get_mask_random_numpy(mask_shape, sparsity, random_state=None):
# def get_mask_random(mask, sparsity, dtype, random_state=None):
# def get_sparsities_erdos_renyi(all_masks,
#                                default_sparsity,
#                                custom_sparsity_map,
#                                include_kernel,
#                                extract_name_fn=mask_extract_name_fn,
#                                erk_power_scale=DEFAULT_ERK_SCALE):
# def get_sparsities_uniform(all_masks,
#                            default_sparsity,
#                            custom_sparsity_map,
#                            extract_name_fn=mask_extract_name_fn):
# def get_sparsities_str(all_masks, default_sparsity):
# def get_sparsities(all_masks,
#                    method,
#                    default_sparsity,
#                    custom_sparsity_map,
#                    extract_name_fn=mask_extract_name_fn,
#                    erk_power_scale=DEFAULT_ERK_SCALE):
# def get_mask_init_fn(all_masks,
#                      method,
#                      default_sparsity,
#                      custom_sparsity_map,
#                      mask_fn=get_mask_random,
#                      erk_power_scale=DEFAULT_ERK_SCALE,
#                      extract_name_fn=mask_extract_name_fn):
# def _get_kernel(layer):
# def get_stats(masked_layers,
#               default_sparsity=0.8,
#               method='erdos_renyi',
#               custom_sparsities=None,
#               is_debug=False,
#               width=1.,
#               first_layer_name='conv1',
#               last_layer_name='conv_preds',
#               param_size=32,
#               erk_power_scale=DEFAULT_ERK_SCALE):


# Path: rigl/sparse_utils_test.py
import absl.testing.parameterized as parameterized
import numpy as np
import tensorflow.compat.v1 as tf
from rigl import sparse_utils

                 tf.get_variable(shape=(1, 1, 3), name='var3/mask')]
    custom_sparsity = {'var3': 0.8}
    sparsities = sparse_utils.get_sparsities(
        all_masks, 'erdos_renyi', default_sparsity, custom_sparsity)
    self.assertEqual(sparsities[all_masks[2].name], 0.8)

  @parameterized.parameters(((2, 3), (2, 3), 0.5),
                            ((1, 1, 2, 3), (1, 1, 2, 3), 0.3),
                            ((8, 6), (4, 3), 0.7),
                            ((80, 4), (20, 20), 0.8),
                            ((2, 6), (2, 3), 0.8))
  def testSparsityDictErdosRenyiSparsitiesScale(
      self, shape1, shape2, default_sparsity):
    _ = self._setup_session()
    all_masks = [tf.get_variable(shape=shape1, name='var1/mask'),
                 tf.get_variable(shape=shape2, name='var2/mask')]
    custom_sparsity = {}
    sparsities = sparse_utils.get_sparsities(
        all_masks, 'erdos_renyi', default_sparsity, custom_sparsity)
    sparsity1 = sparsities[all_masks[0].name]
    size1 = np.prod(shape1)
    sparsity2 = sparsities[all_masks[1].name]
    size2 = np.prod(shape2)
    # Ensure that total number of connections are similar.
    expected_zeros_uniform = (
        sparse_utils.get_n_zeros(size1, default_sparsity) +
        sparse_utils.get_n_zeros(size2, default_sparsity))
    # Ensure that total number of connections are similar.
    expected_zeros_current = (
        sparse_utils.get_n_zeros(size1, sparsity1) +

        sparse_utils.get_n_zeros(size2, sparsity2))