# 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
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# 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.

"""Tests for the data_helper input pipeline and the training process.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function




class GetMaskRandomTest(tf.test.TestCase, parameterized.TestCase):

  def _setup_session(self):
    """Resets the graph and returns a fresh session."""
    tf.reset_default_graph()
    sess = tf.Session()
    return sess

  @parameterized.parameters(((30, 40), 0.5), ((1, 2, 1, 4), 0.8), ((3,), 0.1))
  def testMaskConnectionDeterminism(self, shape, sparsity):
    sess = self._setup_session()
    mask = tf.ones(shape)

# 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

#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# 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.

"""Tests for the data_helper input pipeline and the training process.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function




class GetMaskRandomTest(tf.test.TestCase, parameterized.TestCase):

  def _setup_session(self):
    """Resets the graph and returns a fresh session."""
    tf.reset_default_graph()
    sess = tf.Session()
    return sess

  @parameterized.parameters(((30, 40), 0.5), ((1, 2, 1, 4), 0.8), ((3,), 0.1))
  def testMaskConnectionDeterminism(self, shape, sparsity):
    sess = self._setup_session()
    mask = tf.ones(shape)

    mask1 = sparse_utils.get_mask_random(mask, sparsity, tf.int32)