# 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 re
import numpy as np
import six
from rigl import sparse_utils
from tensorflow.contrib.model_pruning.python import pruning
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import init_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn_ops
from tensorflow.python.ops import random_ops
from tensorflow.python.ops import state_ops
from tensorflow.python.ops import stateless_random_ops
from tensorflow.python.ops import variable_scope
from tensorflow.python.ops import variables
from tensorflow.python.tpu.ops import tpu_ops
from tensorflow.python.training import learning_rate_decay
from tensorflow.python.training import moving_averages
from tensorflow.python.training import optimizer as tf_optimizer
from tensorflow.python.training import training_util
      name: Optional name for the returned operation.  Default to the name
        passed to the `Optimizer` constructor.

    Returns:
      An `Operation` that applies the specified gradients. If `global_step`
      was not None, that operation also increments `global_step`.
    """
    optimizer_update = self._optimizer.apply_gradients(
        grads_and_vars, global_step=global_step, name=name)
    vars_dict = {
        re.findall('(.+)/weights:0', var.name)[0]: var
        for var in self.get_weights()
    }

    def dnw_fn(mask, sparsity, dtype):
      """Creates a mask with smallest magnitudes with deterministic sparsity.

      Args:
        mask: tf.Tensor, used to obtain correct corresponding gradient.
        sparsity: float, between 0 and 1.
        dtype: tf.dtype, type of the return value.

      Returns:
        tf.Tensor
      """
      del dtype
      var_name = sparse_utils.mask_extract_name_fn(mask.name)
      v = vars_dict[var_name]
      score_drop = math_ops.abs(v)
      n_total = np.prod(score_drop.shape.as_list())

# 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_optimizers.py
import re
import numpy as np
import six
from rigl import sparse_utils
from tensorflow.contrib.model_pruning.python import pruning
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import init_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn_ops
from tensorflow.python.ops import random_ops
from tensorflow.python.ops import state_ops
from tensorflow.python.ops import stateless_random_ops
from tensorflow.python.ops import variable_scope
from tensorflow.python.ops import variables
from tensorflow.python.tpu.ops import tpu_ops
from tensorflow.python.training import learning_rate_decay
from tensorflow.python.training import moving_averages
from tensorflow.python.training import optimizer as tf_optimizer
from tensorflow.python.training import training_util

      name: Optional name for the returned operation.  Default to the name
        passed to the `Optimizer` constructor.

    Returns:
      An `Operation` that applies the specified gradients. If `global_step`
      was not None, that operation also increments `global_step`.
    """
    optimizer_update = self._optimizer.apply_gradients(
        grads_and_vars, global_step=global_step, name=name)
    vars_dict = {
        re.findall('(.+)/weights:0', var.name)[0]: var
        for var in self.get_weights()
    }

    def dnw_fn(mask, sparsity, dtype):
      """Creates a mask with smallest magnitudes with deterministic sparsity.

      Args:
        mask: tf.Tensor, used to obtain correct corresponding gradient.
        sparsity: float, between 0 and 1.
        dtype: tf.dtype, type of the return value.

      Returns:
        tf.Tensor
      """
      del dtype
      var_name = sparse_utils.mask_extract_name_fn(mask.name)
      v = vars_dict[var_name]
      score_drop = math_ops.abs(v)
      n_total = np.prod(score_drop.shape.as_list())

      n_prune = sparse_utils.get_n_zeros(n_total, sparsity)