+# tf.internal API namespace++| Status        | Proposed |+:-------------- |:---------------------------------------------------- |+| **RFC #**     | [278](https://github.com/tensorflow/community/pull/278)|+| **Author(s)** | Qianli Zhu ([email protected]) |+| **Sponsor**   | Martin Wicke ([email protected]), Alex Apassos ([email protected])|+| **Updated**   | 2020-08-10                                           |+| **Intended audience**| tf-api-owners, tf-addons, keras-team, deepmind/sonnet|++## Objective++Adding a new "internal" API namespace in TF to host APIs for framework building/testing, etc. The API will have looser contracts compared to core TF API.++## Motivation++TensorFlow has a strong contract to ensure the stability of its API. The core +API, once added, can't be updated with backward incompatible change or removed. +TF API owners will pay extra attention to any new API proposal due to this +restriction, based on the cost we have to bear to maintain the API +compatibility. See more details about API review +[here](https://github.com/tensorflow/community/blob/master/governance/api-reviews.md).++Keras team is trying to split its code into a [standalone Github repository](https://github.com/tensorflow/community/blob/master/rfcs/20200205-standalone-keras-repository.md). +One of the key actions before the repository split is that Keras code has to +rely on the public TF API only, to avoid any future breakages due to private +method/class change.++Historically, Keras was an internal component of TF, and relies heavily on a lot+of private TF utilities and functionalities that are not part of the public API. +Exposing those functions into public APIs will be very confusing to the end +user, since those functions are useful to framework-builders, but not typical +TensorFlow users. In addition, these utilities would be overly restricted by the +standard TensorFlow API contract, as we would like to change/remove them more +often than the existing core API.+++## Design Proposal++Add a new namespace "tensorflow.internal" to host the framework building/testing +related APIs. <b>(See alternative naming in the sections below, the naming of +the namespace will be discussed in the design meeting.)</b>++### Backward compatibility+The contract for all the APIs under this namespace is looser compared to core TF+API. It will remain backward compatible for at least 1 minor TF release. ++This means for any API that is added in 2.x release, it will remain the same in +2.x+1. If we choose to remove the same API, we will mark it as deprecated in +2.x+1, and delete it at 2.x+2 release. TensorFlow is released every 3 months, +and this will give enough time for all the clients to change to the new +API/alternative.++Any deprecation and backward incompatible API change will be explicitly called +out in TF release notes. This applies to both core API and "internal" API.++### Acceptance Criteria+The candidate of the "internal" API should:++1. Does NOT fit for core TF API, otherwise it should be exposed as core TF API.+1. Are effectively required to build/test/deploy a framework or library on top +   of TensorFlow, or to integrate such a framework into TensorFlow's API.+1. Mature enough and won't change behavior/interface for every release. +1. Non-trivial, otherwise the function should be copied to the client side.+1. Widely used in the implementation of TF public APIs (i.e. new functionality +   isn’t immediately added to the tf.internal namespace)+1. Has at least two downstream libraries which are known to need it.++TF API owners will review the new API proposal, and follow the existing review +process for core TF API.+++### Documentation and Docstring+The "internal" API should have the same style and standard as the core +TensorFlow API, which is documented [here](https://github.com/tensorflow/community/blob/master/governance/api-reviews.md#docstrings). +We should explicitly list out the difference between "internal" API and core+API, and also choose a different place on tensorflow.org so that existing user +are not confused.++### Naming and sub namespace+Similar "internal" APIs should be grouped together as sub namespaces, e.g., test +related APIs should live under "tf.internal.test". This is aligned with the +existing TF naming conversion.++Try not to export experimental APIs since the "internal" API should be mature +enough.++### Current candidate+The following list is created from the private TF method usage within Keras, +when we were trying to convert Keras to use the public TF API only. This is by +no means the full list, but will serve as the first wave of review requests we +send to the API owner for review. We don't expect all of them to be approved, +and will discuss with the API owner on a case to case basis.++|Symbol location  |API Name  | +:-------------- |:---------------------------------------------------- |+|python.framework.func_graph.FuncGraph |tf.internal.FuncGraph     |

+# tf.internal API namespace++| Status        | Proposed |+:-------------- |:---------------------------------------------------- |+| **RFC #**     | [278](https://github.com/tensorflow/community/pull/278)|+| **Author(s)** | Qianli Zhu ([email protected]) |+| **Sponsor**   | Martin Wicke ([email protected]), Alex Apassos ([email protected])|+| **Updated**   | 2020-08-10                                           |+| **Intended audience**| tf-api-owners, tf-addons, keras-team, deepmind/sonnet|++## Objective++Adding a new "internal" API namespace in TF to host APIs for framework building/testing, etc. The API will have looser contracts compared to core TF API.++## Motivation++TensorFlow has a strong contract to ensure the stability of its API. The core +API, once added, can't be updated with backward incompatible change or removed. +TF API owners will pay extra attention to any new API proposal due to this +restriction, based on the cost we have to bear to maintain the API +compatibility. See more details about API review +[here](https://github.com/tensorflow/community/blob/master/governance/api-reviews.md).++Keras team is trying to split its code into a [standalone Github repository](https://github.com/tensorflow/community/blob/master/rfcs/20200205-standalone-keras-repository.md). +One of the key actions before the repository split is that Keras code has to +rely on the public TF API only, to avoid any future breakages due to private +method/class change.++Historically, Keras was an internal component of TF, and relies heavily on a lot+of private TF utilities and functionalities that are not part of the public API. +Exposing those functions into public APIs will be very confusing to the end +user, since those functions are useful to framework-builders, but not typical +TensorFlow users. In addition, these utilities would be overly restricted by the +standard TensorFlow API contract, as we would like to change/remove them more +often than the existing core API.+++## Design Proposal++Add a new namespace "tensorflow.internal" to host the framework building/testing +related APIs. <b>(See alternative naming in the sections below, the naming of +the namespace will be discussed in the design meeting.)</b>++### Backward compatibility+The contract for all the APIs under this namespace is looser compared to core TF+API. It will remain backward compatible for at least 1 minor TF release. ++This means for any API that is added in 2.x release, it will remain the same in +2.x+1. If we choose to remove the same API, we will mark it as deprecated in +2.x+1, and delete it at 2.x+2 release. TensorFlow is released every 3 months, +and this will give enough time for all the clients to change to the new +API/alternative.++Any deprecation and backward incompatible API change will be explicitly called +out in TF release notes. This applies to both core API and "internal" API.++### Acceptance Criteria+The candidate of the "internal" API should:++1. Does NOT fit for core TF API, otherwise it should be exposed as core TF API.+1. Are effectively required to build/test/deploy a framework or library on top +   of TensorFlow, or to integrate such a framework into TensorFlow's API.+1. Mature enough and won't change behavior/interface for every release. +1. Non-trivial, otherwise the function should be copied to the client side.+1. Widely used in the implementation of TF public APIs (i.e. new functionality +   isn’t immediately added to the tf.internal namespace)+1. Has at least two downstream libraries which are known to need it.++TF API owners will review the new API proposal, and follow the existing review +process for core TF API.+++### Documentation and Docstring+The "internal" API should have the same style and standard as the core +TensorFlow API, which is documented [here](https://github.com/tensorflow/community/blob/master/governance/api-reviews.md#docstrings). +We should explicitly list out the difference between "internal" API and core+API, and also choose a different place on tensorflow.org so that existing user +are not confused.++### Naming and sub namespace+Similar "internal" APIs should be grouped together as sub namespaces, e.g., test +related APIs should live under "tf.internal.test". This is aligned with the +existing TF naming conversion.++Try not to export experimental APIs since the "internal" API should be mature +enough.++### Current candidate+The following list is created from the private TF method usage within Keras, +when we were trying to convert Keras to use the public TF API only. This is by +no means the full list, but will serve as the first wave of review requests we +send to the API owner for review. We don't expect all of them to be approved, +and will discuss with the API owner on a case to case basis.++|Symbol location  |API Name  | +:-------------- |:---------------------------------------------------- |+|python.framework.func_graph.FuncGraph |tf.internal.FuncGraph     |

+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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 utilities in sparse training pipelines."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++from absl.testing import parameterized+import numpy as np+import tensorflow as tf++from tensorflow.python.keras import keras_parameterized+from tensorflow_model_optimization.python.core.sparsity_tf2 import sparse_utils++dtypes = tf.dtypes+test = tf.test+++def sample_noise(shape=(), mu=0, sigma=1.):+  sample = tf.random.normal(shape, mean=mu,  stddev=sigma, dtype=tf.float64)+  return sample++def calculate_stdev(p):+  return np.sqrt(p * (1 - p))++class SparseUtilsTest(test.TestCase, parameterized.TestCase):++  def setUp(self):+    super(SparseUtilsTest, self).setUp()+    self.ratios = [0.0, 0.1, 0.25, 0.5, 0.75, 0.9, 1.0]+    self.shapes = [(), [1], [5], [2, 2], [4, 4], [8, 4], [3, 17], [3, 6, 9]]+    self.seeds = range(20)++  def matrix_init(self, _shape):+    return tf.ones(_shape)+  +  def matrix_init_with_type(self, _shape, _type):+    tf.ones(_shape, dtype=_type)++  def get_bernouilli(self, p):+    return sparse_utils.Bernouilli(p)++  def get_permuteones(self, ratio):+    return sparse_utils.PermuteOnes(ratio)++  @parameters.parameters(+    (-1.,), (1.5,)+  )+  def testInvalidRatioRaisesError(self, p):+    with self.assertRaises(ValueError):+      bernouilli = self.get_bernouilli(p)+      permuteones = self.get_permuteones(p)++  @parameterized.parameters(+    (i,) for i in self.ratios+  )+  def testScatteredInitialOnes(self, ratio):+    shape = (4, 4)+    matrix = self.matrix_init(shape)+    self.assertAllEqual(shape, matrix.shape)++    def _check_scatter(sparse_matrix):+      ones_indices = np.where(sparse_matrix == 1.)+      mean_nonzeros = np.mean(ones_indices)+      ones_count = np.count_nonzero(sparse_matrix)+      return ones_count, mean_nonzeros++    seed = 0+    ones_indices = []+    ones_counts = []+    for _ in range(100):+      matrix_sparse_permute = self.get_permuteones(ratio)(shape, seed=seed)+      ones_count, mean_nonzeros = _check_scatter(matrix_sparse_permute)+      ones_indices.append(mean_nonzeros)+      ones_counts.append(ones_count)+      seed += 1+    # check that 1/0s are scattered+    # mean of entire matrix and mean of entire nonzeros+    self.assertAllClose(ones_indices[:len(ones_indices)//2], ones_indices[len(ones_indices)//2:])+    self.assertAllClose(ones_counts[:len(ones_counts)//2], ones_counts[len(ones_counts)//2:])++  @parameterized.parameters(+    (shape,) for shape in self.shapes+  )+  def testBernouilliMatrixFraction(self, shape):+    matrix = self.matrix_init(shape)+    self.assertAllEqual(shape, matrix.shape)++    seed = 0+    expected_counts = [16, 1.6, 4, 8, 2, 1.7, 0]+    expected_stdev = list(map(calculate_stdev, ratios))+    counts = []+    for ratio in self.ratios:+      ratio_list = []+      for _ in range(100):+        output = self.get_bernouilli(ratio)(shape, seed=seed)+        ratio_list.append(np.count_nonzero(output))

+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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 utilities in sparse training pipelines."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++from absl.testing import parameterized+import numpy as np+import tensorflow as tf++from tensorflow.python.keras import keras_parameterized+from tensorflow_model_optimization.python.core.sparsity_tf2 import sparse_utils++dtypes = tf.dtypes+test = tf.test+++def sample_noise(shape=(), mu=0, sigma=1.):+  sample = tf.random.normal(shape, mean=mu,  stddev=sigma, dtype=tf.float64)+  return sample++def calculate_stdev(p):+  return np.sqrt(p * (1 - p))++class SparseUtilsTest(test.TestCase, parameterized.TestCase):++  def setUp(self):+    super(SparseUtilsTest, self).setUp()+    self.ratios = [0.0, 0.1, 0.25, 0.5, 0.75, 0.9, 1.0]+    self.shapes = [(), [1], [5], [2, 2], [4, 4], [8, 4], [3, 17], [3, 6, 9]]+    self.seeds = range(20)++  def matrix_init(self, _shape):+    return tf.ones(_shape)+  +  def matrix_init_with_type(self, _shape, _type):+    tf.ones(_shape, dtype=_type)++  def get_bernouilli(self, p):+    return sparse_utils.Bernouilli(p)++  def get_permuteones(self, ratio):+    return sparse_utils.PermuteOnes(ratio)++  @parameters.parameters(+    (-1.,), (1.5,)+  )+  def testInvalidRatioRaisesError(self, p):+    with self.assertRaises(ValueError):+      bernouilli = self.get_bernouilli(p)+      permuteones = self.get_permuteones(p)++  @parameterized.parameters(+    (i,) for i in self.ratios+  )+  def testScatteredInitialOnes(self, ratio):+    shape = (4, 4)+    matrix = self.matrix_init(shape)+    self.assertAllEqual(shape, matrix.shape)++    def _check_scatter(sparse_matrix):+      ones_indices = np.where(sparse_matrix == 1.)+      mean_nonzeros = np.mean(ones_indices)+      ones_count = np.count_nonzero(sparse_matrix)+      return ones_count, mean_nonzeros++    seed = 0+    ones_indices = []+    ones_counts = []+    for _ in range(100):+      matrix_sparse_permute = self.get_permuteones(ratio)(shape, seed=seed)+      ones_count, mean_nonzeros = _check_scatter(matrix_sparse_permute)+      ones_indices.append(mean_nonzeros)+      ones_counts.append(ones_count)+      seed += 1+    # check that 1/0s are scattered+    # mean of entire matrix and mean of entire nonzeros+    self.assertAllClose(ones_indices[:len(ones_indices)//2], ones_indices[len(ones_indices)//2:])+    self.assertAllClose(ones_counts[:len(ones_counts)//2], ones_counts[len(ones_counts)//2:])++  @parameterized.parameters(+    (shape,) for shape in self.shapes+  )+  def testBernouilliMatrixFraction(self, shape):+    matrix = self.matrix_init(shape)+    self.assertAllEqual(shape, matrix.shape)++    seed = 0+    expected_counts = [16, 1.6, 4, 8, 2, 1.7, 0]+    expected_stdev = list(map(calculate_stdev, ratios))+    counts = []+    for ratio in self.ratios:+      ratio_list = []+      for _ in range(100):+        output = self.get_bernouilli(ratio)(shape, seed=seed)+        ratio_list.append(np.count_nonzero(output))+        self.assertAllEqual(output.shape, shape)+        seed += 1+      counts.append(ratio_list)+    mean_counts = np.mean(counts, axis=-1)

+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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.+# ==============================================================================+"""Convenience functions for sparse training."""++class Bernouilli(tf.keras.initializers.Initializer):+  """+  Initialization distributio following a Bernouilli process..+  """++  def __init__(self, p):+    """+    p: probability parameter of success (i.e. 1).+    """+    if not (p >= 0. and p <= 1.):+      raise ValueError('p parameter must be a valid probability, i.e. in [0, 1].')+    self.p = p++  def get_config(self):+    return {'p': self.p}++  @classmethod+  def from_config(cls, config):+    return cls(**config)++  def __call__(self, shape, dytpe=tf.dtypes.float32):+    """Number of zeros = np.ceil(probability * size) in expectation."""+    probs = tf.zeros(shape=list(shape)) + self.p+    uniform = tf.random.uniform(shape)+    initial = tf.less(uniform, probs)++    return tf.cast(initial, dtype=dtype)++class PermuteOnes(tf.keras.initializers.Initializer):+  """+  Initialization of a deterministically sparse matrix.+  This initializer takes in an input ratio and sets exactly+  that ratio of the mask entries as ones  leaving the rest as zeros.+  The ones are detministically, randomly permmuted across the tensor.+  """+  def __init__(self, ratio=None):

+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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 utilities in sparse training pipelines."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++from absl.testing import parameterized+import numpy as np+import tensorflow as tf++from tensorflow.python.keras import keras_parameterized+from tensorflow_model_optimization.python.core.sparsity_tf2 import sparse_utils++dtypes = tf.dtypes+test = tf.test+++def sample_noise(shape=(), mu=0, sigma=1.):+  sample = tf.random.normal(shape, mean=mu,  stddev=sigma, dtype=tf.float64)+  return sample++def calculate_stdev(p):+  return np.sqrt(p * (1 - p))++class SparseUtilsTest(test.TestCase, parameterized.TestCase):++  def setUp(self):+    super(SparseUtilsTest, self).setUp()+    self.ratios = [0.0, 0.1, 0.25, 0.5, 0.75, 0.9, 1.0]+    self.shapes = [(), [1], [5], [2, 2], [4, 4], [8, 4], [3, 17], [3, 6, 9]]+    self.seeds = range(20)++  def matrix_init(self, _shape):+    return tf.ones(_shape)+  +  def matrix_init_with_type(self, _shape, _type):+    tf.ones(_shape, dtype=_type)++  def get_bernouilli(self, p):+    return sparse_utils.Bernouilli(p)++  def get_permuteones(self, ratio):+    return sparse_utils.PermuteOnes(ratio)++  @parameters.parameters(+    (-1.,), (1.5,)+  )+  def testInvalidRatioRaisesError(self, p):+    with self.assertRaises(ValueError):+      bernouilli = self.get_bernouilli(p)+      permuteones = self.get_permuteones(p)++  @parameterized.parameters(+    (i,) for i in self.ratios+  )+  def testScatteredInitialOnes(self, ratio):+    shape = (4, 4)+    matrix = self.matrix_init(shape)+    self.assertAllEqual(shape, matrix.shape)++    def _check_scatter(sparse_matrix):+      ones_indices = np.where(sparse_matrix == 1.)+      mean_nonzeros = np.mean(ones_indices)+      ones_count = np.count_nonzero(sparse_matrix)+      return ones_count, mean_nonzeros++    seed = 0+    ones_indices = []+    ones_counts = []+    for _ in range(100):+      matrix_sparse_permute = self.get_permuteones(ratio)(shape, seed=seed)+      ones_count, mean_nonzeros = _check_scatter(matrix_sparse_permute)+      ones_indices.append(mean_nonzeros)+      ones_counts.append(ones_count)+      seed += 1+    # check that 1/0s are scattered+    # mean of entire matrix and mean of entire nonzeros+    self.assertAllClose(ones_indices[:len(ones_indices)//2], ones_indices[len(ones_indices)//2:])+    self.assertAllClose(ones_counts[:len(ones_counts)//2], ones_counts[len(ones_counts)//2:])++  @parameterized.parameters(+    (shape,) for shape in self.shapes+  )+  def testBernouilliMatrixFraction(self, shape):+    matrix = self.matrix_init(shape)+    self.assertAllEqual(shape, matrix.shape)++    seed = 0+    expected_counts = [16, 1.6, 4, 8, 2, 1.7, 0]

+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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 utilities in sparse training pipelines."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++from absl.testing import parameterized+import numpy as np+import tensorflow as tf++from tensorflow.python.keras import keras_parameterized+from tensorflow_model_optimization.python.core.sparsity_tf2 import sparse_utils++dtypes = tf.dtypes+test = tf.test+++def sample_noise(shape=(), mu=0, sigma=1.):+  sample = tf.random.normal(shape, mean=mu,  stddev=sigma, dtype=tf.float64)+  return sample++def calculate_stdev(p):+  return np.sqrt(p * (1 - p))++class SparseUtilsTest(test.TestCase, parameterized.TestCase):++  def setUp(self):+    super(SparseUtilsTest, self).setUp()+    self.ratios = [0.0, 0.1, 0.25, 0.5, 0.75, 0.9, 1.0]+    self.shapes = [(), [1], [5], [2, 2], [4, 4], [8, 4], [3, 17], [3, 6, 9]]+    self.seeds = range(20)++  def matrix_init(self, _shape):+    return tf.ones(_shape)+  +  def matrix_init_with_type(self, _shape, _type):+    tf.ones(_shape, dtype=_type)++  def get_bernouilli(self, p):+    return sparse_utils.Bernouilli(p)++  def get_permuteones(self, ratio):+    return sparse_utils.PermuteOnes(ratio)++  @parameters.parameters(+    (-1.,), (1.5,)+  )+  def testInvalidRatioRaisesError(self, p):+    with self.assertRaises(ValueError):+      bernouilli = self.get_bernouilli(p)+      permuteones = self.get_permuteones(p)++  @parameterized.parameters(+    (i,) for i in self.ratios+  )+  def testScatteredInitialOnes(self, ratio):+    shape = (4, 4)+    matrix = self.matrix_init(shape)+    self.assertAllEqual(shape, matrix.shape)++    def _check_scatter(sparse_matrix):+      ones_indices = np.where(sparse_matrix == 1.)+      mean_nonzeros = np.mean(ones_indices)+      ones_count = np.count_nonzero(sparse_matrix)+      return ones_count, mean_nonzeros++    seed = 0+    ones_indices = []+    ones_counts = []+    for _ in range(100):+      matrix_sparse_permute = self.get_permuteones(ratio)(shape, seed=seed)+      ones_count, mean_nonzeros = _check_scatter(matrix_sparse_permute)+      ones_indices.append(mean_nonzeros)+      ones_counts.append(ones_count)+      seed += 1+    # check that 1/0s are scattered+    # mean of entire matrix and mean of entire nonzeros+    self.assertAllClose(ones_indices[:len(ones_indices)//2], ones_indices[len(ones_indices)//2:])+    self.assertAllClose(ones_counts[:len(ones_counts)//2], ones_counts[len(ones_counts)//2:])

+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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 utilities in sparse training pipelines."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++from absl.testing import parameterized+import numpy as np+import tensorflow as tf++from tensorflow.python.keras import keras_parameterized+from tensorflow_model_optimization.python.core.sparsity_tf2 import sparse_utils++dtypes = tf.dtypes+test = tf.test+++def sample_noise(shape=(), mu=0, sigma=1.):+  sample = tf.random.normal(shape, mean=mu,  stddev=sigma, dtype=tf.float64)+  return sample++def calculate_stdev(p):+  return np.sqrt(p * (1 - p))++class SparseUtilsTest(test.TestCase, parameterized.TestCase):++  def setUp(self):+    super(SparseUtilsTest, self).setUp()+    self.ratios = [0.0, 0.1, 0.25, 0.5, 0.75, 0.9, 1.0]+    self.shapes = [(), [1], [5], [2, 2], [4, 4], [8, 4], [3, 17], [3, 6, 9]]+    self.seeds = range(20)++  def matrix_init(self, _shape):+    return tf.ones(_shape)+  +  def matrix_init_with_type(self, _shape, _type):+    tf.ones(_shape, dtype=_type)++  def get_bernouilli(self, p):+    return sparse_utils.Bernouilli(p)++  def get_permuteones(self, ratio):+    return sparse_utils.PermuteOnes(ratio)++  @parameters.parameters(+    (-1.,), (1.5,)+  )+  def testInvalidRatioRaisesError(self, p):+    with self.assertRaises(ValueError):+      bernouilli = self.get_bernouilli(p)+      permuteones = self.get_permuteones(p)++  @parameterized.parameters(+    (i,) for i in self.ratios+  )+  def testScatteredInitialOnes(self, ratio):+    shape = (4, 4)+    matrix = self.matrix_init(shape)+    self.assertAllEqual(shape, matrix.shape)++    def _check_scatter(sparse_matrix):+      ones_indices = np.where(sparse_matrix == 1.)+      mean_nonzeros = np.mean(ones_indices)+      ones_count = np.count_nonzero(sparse_matrix)+      return ones_count, mean_nonzeros++    seed = 0+    ones_indices = []+    ones_counts = []+    for _ in range(100):+      matrix_sparse_permute = self.get_permuteones(ratio)(shape, seed=seed)+      ones_count, mean_nonzeros = _check_scatter(matrix_sparse_permute)+      ones_indices.append(mean_nonzeros)+      ones_counts.append(ones_count)+      seed += 1+    # check that 1/0s are scattered+    # mean of entire matrix and mean of entire nonzeros+    self.assertAllClose(ones_indices[:len(ones_indices)//2], ones_indices[len(ones_indices)//2:])+    self.assertAllClose(ones_counts[:len(ones_counts)//2], ones_counts[len(ones_counts)//2:])

+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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 utilities in sparse training pipelines."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++from absl.testing import parameterized+import numpy as np+import tensorflow as tf++from tensorflow.python.keras import keras_parameterized+from tensorflow_model_optimization.python.core.sparsity_tf2 import sparse_utils++dtypes = tf.dtypes+test = tf.test+++def sample_noise(shape=(), mu=0, sigma=1.):+  sample = tf.random.normal(shape, mean=mu,  stddev=sigma, dtype=tf.float64)+  return sample++def calculate_stdev(p):+  return np.sqrt(p * (1 - p))++class SparseUtilsTest(test.TestCase, parameterized.TestCase):++  def setUp(self):+    super(SparseUtilsTest, self).setUp()+    self.ratios = [0.0, 0.1, 0.25, 0.5, 0.75, 0.9, 1.0]+    self.shapes = [(), [1], [5], [2, 2], [4, 4], [8, 4], [3, 17], [3, 6, 9]]+    self.seeds = range(20)++  def matrix_init(self, _shape):+    return tf.ones(_shape)+  +  def matrix_init_with_type(self, _shape, _type):+    tf.ones(_shape, dtype=_type)++  def get_bernouilli(self, p):+    return sparse_utils.Bernouilli(p)++  def get_permuteones(self, ratio):+    return sparse_utils.PermuteOnes(ratio)++  @parameters.parameters(+    (-1.,), (1.5,)+  )+  def testInvalidRatioRaisesError(self, p):+    with self.assertRaises(ValueError):+      bernouilli = self.get_bernouilli(p)+      permuteones = self.get_permuteones(p)++  @parameterized.parameters(+    (i,) for i in self.ratios+  )+  def testScatteredInitialOnes(self, ratio):+    shape = (4, 4)+    matrix = self.matrix_init(shape)+    self.assertAllEqual(shape, matrix.shape)++    def _check_scatter(sparse_matrix):+      ones_indices = np.where(sparse_matrix == 1.)+      mean_nonzeros = np.mean(ones_indices)+      ones_count = np.count_nonzero(sparse_matrix)+      return ones_count, mean_nonzeros++    seed = 0+    ones_indices = []+    ones_counts = []+    for _ in range(100):+      matrix_sparse_permute = self.get_permuteones(ratio)(shape, seed=seed)+      ones_count, mean_nonzeros = _check_scatter(matrix_sparse_permute)+      ones_indices.append(mean_nonzeros)+      ones_counts.append(ones_count)+      seed += 1+    # check that 1/0s are scattered+    # mean of entire matrix and mean of entire nonzeros+    self.assertAllClose(ones_indices[:len(ones_indices)//2], ones_indices[len(ones_indices)//2:])

+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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 utilities in sparse training pipelines."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++from absl.testing import parameterized+import numpy as np+import tensorflow as tf++from tensorflow.python.keras import keras_parameterized+from tensorflow_model_optimization.python.core.sparsity_tf2 import sparse_utils++dtypes = tf.dtypes+test = tf.test+++def sample_noise(shape=(), mu=0, sigma=1.):+  sample = tf.random.normal(shape, mean=mu,  stddev=sigma, dtype=tf.float64)+  return sample++def calculate_stdev(p):+  return np.sqrt(p * (1 - p))++class SparseUtilsTest(test.TestCase, parameterized.TestCase):++  def setUp(self):+    super(SparseUtilsTest, self).setUp()+    self.ratios = [0.0, 0.1, 0.25, 0.5, 0.75, 0.9, 1.0]+    self.shapes = [(), [1], [5], [2, 2], [4, 4], [8, 4], [3, 17], [3, 6, 9]]+    self.seeds = range(20)++  def matrix_init(self, _shape):

+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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 utilities in sparse training pipelines."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++from absl.testing import parameterized+import numpy as np+import tensorflow as tf++from tensorflow.python.keras import keras_parameterized+from tensorflow_model_optimization.python.core.sparsity_tf2 import sparse_utils++dtypes = tf.dtypes+test = tf.test+++def sample_noise(shape=(), mu=0, sigma=1.):+  sample = tf.random.normal(shape, mean=mu,  stddev=sigma, dtype=tf.float64)+  return sample++def calculate_stdev(p):+  return np.sqrt(p * (1 - p))++class SparseUtilsTest(test.TestCase, parameterized.TestCase):++  def setUp(self):+    super(SparseUtilsTest, self).setUp()+    self.ratios = [0.0, 0.1, 0.25, 0.5, 0.75, 0.9, 1.0]+    self.shapes = [(), [1], [5], [2, 2], [4, 4], [8, 4], [3, 17], [3, 6, 9]]+    self.seeds = range(20)++  def matrix_init(self, _shape):+    return tf.ones(_shape)+  +  def matrix_init_with_type(self, _shape, _type):

+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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 utilities in sparse training pipelines."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++from absl.testing import parameterized+import numpy as np+import tensorflow as tf++from tensorflow.python.keras import keras_parameterized+from tensorflow_model_optimization.python.core.sparsity_tf2 import sparse_utils++dtypes = tf.dtypes+test = tf.test+++def sample_noise(shape=(), mu=0, sigma=1.):+  sample = tf.random.normal(shape, mean=mu,  stddev=sigma, dtype=tf.float64)+  return sample++def calculate_stdev(p):+  return np.sqrt(p * (1 - p))++class SparseUtilsTest(test.TestCase, parameterized.TestCase):++  def setUp(self):+    super(SparseUtilsTest, self).setUp()+    self.ratios = [0.0, 0.1, 0.25, 0.5, 0.75, 0.9, 1.0]+    self.shapes = [(), [1], [5], [2, 2], [4, 4], [8, 4], [3, 17], [3, 6, 9]]+    self.seeds = range(20)++  def matrix_init(self, _shape):+    return tf.ones(_shape)+  +  def matrix_init_with_type(self, _shape, _type):+    tf.ones(_shape, dtype=_type)++  def get_bernouilli(self, p):+    return sparse_utils.Bernouilli(p)++  def get_permuteones(self, ratio):+    return sparse_utils.PermuteOnes(ratio)++  @parameters.parameters(+    (-1.,), (1.5,)+  )+  def testInvalidRatioRaisesError(self, p):+    with self.assertRaises(ValueError):+      bernouilli = self.get_bernouilli(p)+      permuteones = self.get_permuteones(p)++  @parameterized.parameters(+    (i,) for i in self.ratios+  )+  def testScatteredInitialOnes(self, ratio):+    shape = (4, 4)+    matrix = self.matrix_init(shape)+    self.assertAllEqual(shape, matrix.shape)++    def _check_scatter(sparse_matrix):+      ones_indices = np.where(sparse_matrix == 1.)+      mean_nonzeros = np.mean(ones_indices)

+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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 utilities in sparse training pipelines."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++from absl.testing import parameterized+import numpy as np+import tensorflow as tf++from tensorflow.python.keras import keras_parameterized+from tensorflow_model_optimization.python.core.sparsity_tf2 import sparse_utils++dtypes = tf.dtypes+test = tf.test+++def sample_noise(shape=(), mu=0, sigma=1.):+  sample = tf.random.normal(shape, mean=mu,  stddev=sigma, dtype=tf.float64)+  return sample++def calculate_stdev(p):+  return np.sqrt(p * (1 - p))++class SparseUtilsTest(test.TestCase, parameterized.TestCase):++  def setUp(self):+    super(SparseUtilsTest, self).setUp()+    self.ratios = [0.0, 0.1, 0.25, 0.5, 0.75, 0.9, 1.0]+    self.shapes = [(), [1], [5], [2, 2], [4, 4], [8, 4], [3, 17], [3, 6, 9]]+    self.seeds = range(20)++  def matrix_init(self, _shape):+    return tf.ones(_shape)+  +  def matrix_init_with_type(self, _shape, _type):+    tf.ones(_shape, dtype=_type)++  def get_bernouilli(self, p):+    return sparse_utils.Bernouilli(p)++  def get_permuteones(self, ratio):

+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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 utilities in sparse training pipelines."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++from absl.testing import parameterized+import numpy as np+import tensorflow as tf++from tensorflow.python.keras import keras_parameterized+from tensorflow_model_optimization.python.core.sparsity_tf2 import sparse_utils++dtypes = tf.dtypes+test = tf.test+++def sample_noise(shape=(), mu=0, sigma=1.):+  sample = tf.random.normal(shape, mean=mu,  stddev=sigma, dtype=tf.float64)+  return sample++def calculate_stdev(p):+  return np.sqrt(p * (1 - p))++class SparseUtilsTest(test.TestCase, parameterized.TestCase):++  def setUp(self):+    super(SparseUtilsTest, self).setUp()+    self.ratios = [0.0, 0.1, 0.25, 0.5, 0.75, 0.9, 1.0]+    self.shapes = [(), [1], [5], [2, 2], [4, 4], [8, 4], [3, 17], [3, 6, 9]]+    self.seeds = range(20)++  def matrix_init(self, _shape):+    return tf.ones(_shape)+  +  def matrix_init_with_type(self, _shape, _type):+    tf.ones(_shape, dtype=_type)++  def get_bernouilli(self, p):

+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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 utilities in sparse training pipelines."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++from absl.testing import parameterized+import numpy as np+import tensorflow as tf++from tensorflow.python.keras import keras_parameterized+from tensorflow_model_optimization.python.core.sparsity_tf2 import sparse_utils++dtypes = tf.dtypes+test = tf.test+++def sample_noise(shape=(), mu=0, sigma=1.):+  sample = tf.random.normal(shape, mean=mu,  stddev=sigma, dtype=tf.float64)+  return sample++def calculate_stdev(p):+  return np.sqrt(p * (1 - p))++class SparseUtilsTest(test.TestCase, parameterized.TestCase):++  def setUp(self):+    super(SparseUtilsTest, self).setUp()+    self.ratios = [0.0, 0.1, 0.25, 0.5, 0.75, 0.9, 1.0]+    self.shapes = [(), [1], [5], [2, 2], [4, 4], [8, 4], [3, 17], [3, 6, 9]]+    self.seeds = range(20)++  def matrix_init(self, _shape):+    return tf.ones(_shape)+  +  def matrix_init_with_type(self, _shape, _type):+    tf.ones(_shape, dtype=_type)++  def get_bernouilli(self, p):+    return sparse_utils.Bernouilli(p)++  def get_permuteones(self, ratio):+    return sparse_utils.PermuteOnes(ratio)++  @parameters.parameters(+    (-1.,), (1.5,)+  )+  def testInvalidRatioRaisesError(self, p):+    with self.assertRaises(ValueError):+      bernouilli = self.get_bernouilli(p)+      permuteones = self.get_permuteones(p)++  @parameterized.parameters(+    (i,) for i in self.ratios

+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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.+# ==============================================================================+"""Convenience functions for sparse training."""++class Bernouilli(tf.keras.initializers.Initializer):+  """+  Initialization distributio following a Bernouilli process..+  """++  def __init__(self, p):+    """+    p: probability parameter of success (i.e. 1).+    """+    if not (p >= 0. and p <= 1.):+      raise ValueError('p parameter must be a valid probability, i.e. in [0, 1].')+    self.p = p++  def get_config(self):+    return {'p': self.p}++  @classmethod+  def from_config(cls, config):+    return cls(**config)++  def __call__(self, shape, dytpe=tf.dtypes.float32):+    """Number of zeros = np.ceil(probability * size) in expectation."""+    probs = tf.zeros(shape=list(shape)) + self.p+    uniform = tf.random.uniform(shape)+    initial = tf.less(uniform, probs)++    return tf.cast(initial, dtype=dtype)++class PermuteOnes(tf.keras.initializers.Initializer):+  """+  Initialization of a deterministically sparse matrix.+  This initializer takes in an input ratio and sets exactly+  that ratio of the mask entries as ones  leaving the rest as zeros.+  The ones are detministically, randomly permmuted across the tensor.+  """+  def __init__(self, ratio=None):+    """+    ratio: the exact number of 1s sampled.+    If ratio is None, will sample randomly from uniform distribution for sparsity.+    """+    if ratio is not None and not (ratio >= 0. and ratio <= 1.):+      raise ValueError('ratio parameter must be a valid percentage, i.e. in [0, 1].')+    self.ratio = ratio if ratio else tf.random.uniform(())

+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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.+# ==============================================================================+"""Convenience functions for sparse training."""++class Bernouilli(tf.keras.initializers.Initializer):+  """+  Initialization distributio following a Bernouilli process..+  """++  def __init__(self, p):+    """+    p: probability parameter of success (i.e. 1).+    """+    if not (p >= 0. and p <= 1.):+      raise ValueError('p parameter must be a valid probability, i.e. in [0, 1].')+    self.p = p++  def get_config(self):+    return {'p': self.p}++  @classmethod+  def from_config(cls, config):+    return cls(**config)++  def __call__(self, shape, dytpe=tf.dtypes.float32):+    """Number of zeros = np.ceil(probability * size) in expectation."""+    probs = tf.zeros(shape=list(shape)) + self.p+    uniform = tf.random.uniform(shape)+    initial = tf.less(uniform, probs)++    return tf.cast(initial, dtype=dtype)++class PermuteOnes(tf.keras.initializers.Initializer):+  """+  Initialization of a deterministically sparse matrix.+  This initializer takes in an input ratio and sets exactly+  that ratio of the mask entries as ones  leaving the rest as zeros.+  The ones are detministically, randomly permmuted across the tensor.

+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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.+# ==============================================================================+"""Convenience functions for sparse training."""++class Bernouilli(tf.keras.initializers.Initializer):+  """+  Initialization distributio following a Bernouilli process..+  """++  def __init__(self, p):+    """+    p: probability parameter of success (i.e. 1).+    """+    if not (p >= 0. and p <= 1.):+      raise ValueError('p parameter must be a valid probability, i.e. in [0, 1].')+    self.p = p++  def get_config(self):+    return {'p': self.p}++  @classmethod+  def from_config(cls, config):+    return cls(**config)++  def __call__(self, shape, dytpe=tf.dtypes.float32):+    """Number of zeros = np.ceil(probability * size) in expectation."""+    probs = tf.zeros(shape=list(shape)) + self.p+    uniform = tf.random.uniform(shape)+    initial = tf.less(uniform, probs)++    return tf.cast(initial, dtype=dtype)++class PermuteOnes(tf.keras.initializers.Initializer):+  """+  Initialization of a deterministically sparse matrix.+  This initializer takes in an input ratio and sets exactly+  that ratio of the mask entries as ones  leaving the rest as zeros.+  The ones are detministically, randomly permmuted across the tensor.+  """+  def __init__(self, ratio=None):

 py_library(     ], ) +py_library(+    name = "sparse_utils",+    srcs = ["sparse_utils.py"],+    srcs_version = "PY2AND3",+    visibility = ["//visibility:public"],+    deps = [+        # tensorflow dep1,+        # python:summary tensorflow dep2,

+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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 utilities in sparse training pipelines."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++from absl.testing import parameterized+import numpy as np+import tensorflow as tf++from tensorflow.python.keras import keras_parameterized+from tensorflow_model_optimization.python.core.sparsity.keras import pruning_schedule+from tensorflow_model_optimization.python.core.sparsity_tf2 import lthpruner as pruner+from tensorflow_model_optimization.python.core.sparsity_tf2 import sparse_utils++dtypes = tf.dtypes+test = tf.test+++def sample_noise(x, mu=0, sigma=1.):+  sample = tf.random.normal((), mean=mu,  stddev=sigma, dtype=tf.float64)+  return sample++class SparseUtilsTest(test.TestCase, parameterized.TestCase):++  def setUp(self):+    super(SparseUtilsTest, self).setUp()+    self.mask_init = lambda x: tf.ones(x)+    self.mask_init_w_type = lambda x, y: tf.ones(x, dtype=y)+    self.bernouilli = lambda p: sparse_utils.Bernouilli(p)+    self.permuteones = lambda ratio: sparse_utils.PermuteOnes(ratio)++  def testBernouilliMaskFraction(self):+    shape = (4, 4)+    mask = self.mask_init(shape)+    self.assertAllEqual(shape, mask.shape)++    seed = 0+    ratio = 0.5+    counts = []+    for _ in range(100):+      tmp = self.bernouilli(ratio)(mask.shape, seed=0)+      counts.append(np.count_nonzero(tmp))+    mean_count = np.mean(counts)+    self.assertAllEqual(mean_count, 4)++  def testDeterministicMaskFraction(self):+    # PermuteOnes+    shape = (4, 4)+    mask = self.mask_init(shape)+    self.assertAllEqual(shape, mask.shape)++    ratio = 0.5+    mask_sparse = self.permuteones(ratio)(mask.shape)+    self.assertAllEqual(np.count_nonzero(mask_sparse), 4)++  def testMaskDeterminism(self):+    shape = (4, 4)+    mask1 = self.mask_init(shape)+    mask2 = self.mask_init(shape)+    self.assertAllEqual(shape, mask1.shape)+    self.assertAllEqual(shape, mask2.shape)++    ratio = 0.5+    seed = 0+    permuteones = self.permuteones(ratio)+    mask1_sparse = permuteones(shape=mask1.shape, seed=seed) +    mask2_sparse = permuteones(shape=mask2.shape, seed=seed)+    self.assertAllEqual(mask1_sparse, mask2_sparse)++  def testMaskDtype(self):+    dtypes = [tf.int32, tf.float32, tf.int64, tf.float64]

+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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 utilities in sparse training pipelines."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++from absl.testing import parameterized+import numpy as np+import tensorflow as tf++from tensorflow.python.keras import keras_parameterized+from tensorflow_model_optimization.python.core.sparsity.keras import pruning_schedule+from tensorflow_model_optimization.python.core.sparsity_tf2 import lthpruner as pruner+from tensorflow_model_optimization.python.core.sparsity_tf2 import sparse_utils++dtypes = tf.dtypes+test = tf.test+++def sample_noise(x, mu=0, sigma=1.):+  sample = tf.random.normal((), mean=mu,  stddev=sigma, dtype=tf.float64)+  return sample++class SparseUtilsTest(test.TestCase, parameterized.TestCase):++  def setUp(self):+    super(SparseUtilsTest, self).setUp()+    self.mask_init = lambda x: tf.ones(x)+    self.mask_init_w_type = lambda x, y: tf.ones(x, dtype=y)+    self.bernouilli = lambda p: sparse_utils.Bernouilli(p)+    self.permuteones = lambda ratio: sparse_utils.PermuteOnes(ratio)++  def testBernouilliMaskFraction(self):+    shape = (4, 4)+    mask = self.mask_init(shape)+    self.assertAllEqual(shape, mask.shape)++    seed = 0+    ratio = 0.5+    counts = []+    for _ in range(100):+      tmp = self.bernouilli(ratio)(mask.shape, seed=0)+      counts.append(np.count_nonzero(tmp))

+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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 utilities in sparse training pipelines."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++from absl.testing import parameterized+import numpy as np+import tensorflow as tf++from tensorflow.python.keras import keras_parameterized+from tensorflow_model_optimization.python.core.sparsity.keras import pruning_schedule+from tensorflow_model_optimization.python.core.sparsity_tf2 import lthpruner as pruner+from tensorflow_model_optimization.python.core.sparsity_tf2 import sparse_utils++dtypes = tf.dtypes+test = tf.test+++def sample_noise(x, mu=0, sigma=1.):+  sample = tf.random.normal((), mean=mu,  stddev=sigma, dtype=tf.float64)+  return sample++class SparseUtilsTest(test.TestCase, parameterized.TestCase):++  def setUp(self):+    super(SparseUtilsTest, self).setUp()+    self.mask_init = lambda x: tf.ones(x)+    self.mask_init_w_type = lambda x, y: tf.ones(x, dtype=y)+    self.bernouilli = lambda p: sparse_utils.Bernouilli(p)+    self.permuteones = lambda ratio: sparse_utils.PermuteOnes(ratio)++  def testBernouilliMaskFraction(self):+    shape = (4, 4)+    mask = self.mask_init(shape)+    self.assertAllEqual(shape, mask.shape)++    seed = 0+    ratio = 0.5+    counts = []+    for _ in range(100):+      tmp = self.bernouilli(ratio)(mask.shape, seed=0)+      counts.append(np.count_nonzero(tmp))+    mean_count = np.mean(counts)+    self.assertAllEqual(mean_count, 4)++  def testDeterministicMaskFraction(self):+    # PermuteOnes+    shape = (4, 4)+    mask = self.mask_init(shape)+    self.assertAllEqual(shape, mask.shape)++    ratio = 0.5+    mask_sparse = self.permuteones(ratio)(mask.shape)+    self.assertAllEqual(np.count_nonzero(mask_sparse), 4)++  def testMaskDeterminism(self):

+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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 utilities in sparse training pipelines."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++from absl.testing import parameterized+import numpy as np+import tensorflow as tf++from tensorflow.python.keras import keras_parameterized+from tensorflow_model_optimization.python.core.sparsity.keras import pruning_schedule+from tensorflow_model_optimization.python.core.sparsity_tf2 import lthpruner as pruner+from tensorflow_model_optimization.python.core.sparsity_tf2 import sparse_utils++dtypes = tf.dtypes+test = tf.test+++def sample_noise(x, mu=0, sigma=1.):+  sample = tf.random.normal((), mean=mu,  stddev=sigma, dtype=tf.float64)+  return sample++class SparseUtilsTest(test.TestCase, parameterized.TestCase):++  def setUp(self):+    super(SparseUtilsTest, self).setUp()+    self.mask_init = lambda x: tf.ones(x)+    self.mask_init_w_type = lambda x, y: tf.ones(x, dtype=y)+    self.bernouilli = lambda p: sparse_utils.Bernouilli(p)+    self.permuteones = lambda ratio: sparse_utils.PermuteOnes(ratio)++  def testBernouilliMaskFraction(self):+    shape = (4, 4)+    mask = self.mask_init(shape)+    self.assertAllEqual(shape, mask.shape)++    seed = 0+    ratio = 0.5+    counts = []+    for _ in range(100):+      tmp = self.bernouilli(ratio)(mask.shape, seed=0)+      counts.append(np.count_nonzero(tmp))+    mean_count = np.mean(counts)+    self.assertAllEqual(mean_count, 4)++  def testDeterministicMaskFraction(self):

+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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 utilities in sparse training pipelines."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++from absl.testing import parameterized+import numpy as np+import tensorflow as tf++from tensorflow.python.keras import keras_parameterized+from tensorflow_model_optimization.python.core.sparsity.keras import pruning_schedule+from tensorflow_model_optimization.python.core.sparsity_tf2 import lthpruner as pruner+from tensorflow_model_optimization.python.core.sparsity_tf2 import sparse_utils++dtypes = tf.dtypes+test = tf.test+++def sample_noise(x, mu=0, sigma=1.):+  sample = tf.random.normal((), mean=mu,  stddev=sigma, dtype=tf.float64)+  return sample++class SparseUtilsTest(test.TestCase, parameterized.TestCase):++  def setUp(self):+    super(SparseUtilsTest, self).setUp()+    self.mask_init = lambda x: tf.ones(x)+    self.mask_init_w_type = lambda x, y: tf.ones(x, dtype=y)+    self.bernouilli = lambda p: sparse_utils.Bernouilli(p)+    self.permuteones = lambda ratio: sparse_utils.PermuteOnes(ratio)++  def testBernouilliMaskFraction(self):+    shape = (4, 4)+    mask = self.mask_init(shape)+    self.assertAllEqual(shape, mask.shape)++    seed = 0+    ratio = 0.5+    counts = []+    for _ in range(100):+      tmp = self.bernouilli(ratio)(mask.shape, seed=0)+      counts.append(np.count_nonzero(tmp))+    mean_count = np.mean(counts)+    self.assertAllEqual(mean_count, 4)

+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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 utilities in sparse training pipelines."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++from absl.testing import parameterized+import numpy as np+import tensorflow as tf++from tensorflow.python.keras import keras_parameterized+from tensorflow_model_optimization.python.core.sparsity.keras import pruning_schedule+from tensorflow_model_optimization.python.core.sparsity_tf2 import lthpruner as pruner+from tensorflow_model_optimization.python.core.sparsity_tf2 import sparse_utils++dtypes = tf.dtypes+test = tf.test+++def sample_noise(x, mu=0, sigma=1.):+  sample = tf.random.normal((), mean=mu,  stddev=sigma, dtype=tf.float64)+  return sample++class SparseUtilsTest(test.TestCase, parameterized.TestCase):++  def setUp(self):+    super(SparseUtilsTest, self).setUp()+    self.mask_init = lambda x: tf.ones(x)+    self.mask_init_w_type = lambda x, y: tf.ones(x, dtype=y)+    self.bernouilli = lambda p: sparse_utils.Bernouilli(p)+    self.permuteones = lambda ratio: sparse_utils.PermuteOnes(ratio)++  def testBernouilliMaskFraction(self):+    shape = (4, 4)+    mask = self.mask_init(shape)+    self.assertAllEqual(shape, mask.shape)++    seed = 0+    ratio = 0.5+    counts = []+    for _ in range(100):+      tmp = self.bernouilli(ratio)(mask.shape, seed=0)+      counts.append(np.count_nonzero(tmp))+    mean_count = np.mean(counts)+    self.assertAllEqual(mean_count, 4)

+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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 utilities in sparse training pipelines."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++from absl.testing import parameterized+import numpy as np+import tensorflow as tf++from tensorflow.python.keras import keras_parameterized+from tensorflow_model_optimization.python.core.sparsity.keras import pruning_schedule+from tensorflow_model_optimization.python.core.sparsity_tf2 import lthpruner as pruner+from tensorflow_model_optimization.python.core.sparsity_tf2 import sparse_utils++dtypes = tf.dtypes+test = tf.test+++def sample_noise(x, mu=0, sigma=1.):+  sample = tf.random.normal((), mean=mu,  stddev=sigma, dtype=tf.float64)+  return sample++class SparseUtilsTest(test.TestCase, parameterized.TestCase):++  def setUp(self):+    super(SparseUtilsTest, self).setUp()+    self.mask_init = lambda x: tf.ones(x)+    self.mask_init_w_type = lambda x, y: tf.ones(x, dtype=y)+    self.bernouilli = lambda p: sparse_utils.Bernouilli(p)+    self.permuteones = lambda ratio: sparse_utils.PermuteOnes(ratio)++  def testBernouilliMaskFraction(self):+    shape = (4, 4)+    mask = self.mask_init(shape)+    self.assertAllEqual(shape, mask.shape)++    seed = 0+    ratio = 0.5+    counts = []+    for _ in range(100):+      tmp = self.bernouilli(ratio)(mask.shape, seed=0)

+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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 utilities in sparse training pipelines."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++from absl.testing import parameterized+import numpy as np+import tensorflow as tf++from tensorflow.python.keras import keras_parameterized+from tensorflow_model_optimization.python.core.sparsity.keras import pruning_schedule+from tensorflow_model_optimization.python.core.sparsity_tf2 import lthpruner as pruner+from tensorflow_model_optimization.python.core.sparsity_tf2 import sparse_utils++dtypes = tf.dtypes+test = tf.test+++def sample_noise(x, mu=0, sigma=1.):+  sample = tf.random.normal((), mean=mu,  stddev=sigma, dtype=tf.float64)+  return sample++class SparseUtilsTest(test.TestCase, parameterized.TestCase):++  def setUp(self):+    super(SparseUtilsTest, self).setUp()+    self.mask_init = lambda x: tf.ones(x)+    self.mask_init_w_type = lambda x, y: tf.ones(x, dtype=y)+    self.bernouilli = lambda p: sparse_utils.Bernouilli(p)+    self.permuteones = lambda ratio: sparse_utils.PermuteOnes(ratio)++  def testBernouilliMaskFraction(self):+    shape = (4, 4)+    mask = self.mask_init(shape)

+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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.+# ==============================================================================+"""Convenience functions for sparse training."""++class Bernouilli(tf.keras.initializers.Initializer):+  """+  Initialization distributio following a Bernouilli process..+  """++  def __init__(self, p):+    """+    p: probability parameter of success (i.e. 1).+    """+    self.p = p++  def get_config(self):+    return {'p': self.p}++  @classmethod+  def from_config(cls, config):+    return cls(**config)++  def __call__(self, shape, dytpe=tf.dtypes.float32):+    """Number of zeros = np.ceil(sparsity * size) in expectation."""

+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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 utilities in sparse training pipelines."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++from absl.testing import parameterized+import numpy as np+import tensorflow as tf++from tensorflow.python.keras import keras_parameterized+from tensorflow_model_optimization.python.core.sparsity.keras import pruning_schedule+from tensorflow_model_optimization.python.core.sparsity_tf2 import lthpruner as pruner+from tensorflow_model_optimization.python.core.sparsity_tf2 import sparse_utils++dtypes = tf.dtypes+test = tf.test+++def sample_noise(x, mu=0, sigma=1.):+  sample = tf.random.normal((), mean=mu,  stddev=sigma, dtype=tf.float64)+  return sample++class SparseUtilsTest(test.TestCase, parameterized.TestCase):++  def setUp(self):+    super(SparseUtilsTest, self).setUp()+    self.mask_init = lambda x: tf.ones(x)

+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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 utilities in sparse training pipelines."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++from absl.testing import parameterized+import numpy as np+import tensorflow as tf++from tensorflow.python.keras import keras_parameterized+from tensorflow_model_optimization.python.core.sparsity.keras import pruning_schedule+from tensorflow_model_optimization.python.core.sparsity_tf2 import lthpruner as pruner+from tensorflow_model_optimization.python.core.sparsity_tf2 import sparse_utils++dtypes = tf.dtypes+test = tf.test+++def sample_noise(x, mu=0, sigma=1.):

+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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 utilities in sparse training pipelines."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++from absl.testing import parameterized+import numpy as np+import tensorflow as tf++from tensorflow.python.keras import keras_parameterized+from tensorflow_model_optimization.python.core.sparsity.keras import pruning_schedule+from tensorflow_model_optimization.python.core.sparsity_tf2 import lthpruner as pruner+from tensorflow_model_optimization.python.core.sparsity_tf2 import sparse_utils++dtypes = tf.dtypes+test = tf.test++def make_pruning_schedule(target_sparsity, begin, end, freq):+  return pruning_schedule.ConstantSparsity(target_sparsity, begin, end, freq)++def sample_noise(x, mu=0, sigma=1.):+  sample = tf.random.normal((), mean=mu,  stddev=sigma, dtype=tf.float64)+  return sample++def _dummy_gradient(x, dtype=tf.float32):+  try:+    base_type = x.dtype+  except:+    base_type = dtype+  grad = tf.ones_like(x, dtype=base_type)+  return grad++class SparseUtilsTest(test.TestCase, parameterized.TestCase):

+class Bernouilli(tf.keras.initializers.Initializer):+  """+  Initialization distributio following a Bernouilli process..+  """++  def __init__(self, p):+    """+    p: probability parameter of success (i.e. 1).+    """+    self.p = p++  def get_config(self):+    return {'p': self.p}++  @classmethod+  def from_config(cls, config):+    return cls(**config)++  def __call__(self, shape, dytpe=tf.dtypes.float32):+    """Number of zeros = np.ceil(sparsity * size) in expectation."""+    probs = tf.zeros(shape=list(shape)) + self.p+    uniform = tf.random.uniform(shape)+    initial = tf.less(uniform, probs)++    return tf.cast(initial, dtype=dtype)++class PermuteOnes(tf.keras.initializers.Initializer):+  """+  Initialization of a deterministically sparse matrix.+  """+  def __init__(self, p=None):+    """+    p: probability parameter of success (i.e. 1).+    If p is None, will sample randomly from uniform distribution for sparsity.+    """+    self.p = p if p else tf.random.uniform(())++  def get_n_ones(self, shape, dtype=tf.dtypes.float32):+    sparsity = self.p if self.p else 0.0+    return tf.math.ceil(sparsity * tf.cast(tf.math.reduce_sum(shape)), dtype=dtype)++  def __call__(self, shape, dtype=tf.dtypes.float32, seed=None):+    flat_mask = tf.reshape(tf.ones(shape), (-1,))+    num_ones = self.get_n_ones(shape, dtype)+    _indices = tf.cast(tf.reshape(tf.linspace(0, num_ones - 1, int(num_ones)), (-1,)), tf.int32)+    indices = tf.reshape(_indices, (-1, 1))+    updates = tf.ones_like(_indices)+    flat_shape = flat_mask.shape+    unshuffled_mask = tf.scatter_nd(indices, udpates, flat_shape)+    shuffled_mask = tf.random.shuffle(unshuffled_mas, seed=seed)++    return tf.reshape(shuffled_mask, shape)+++class ErdosRenyi(tf.keras.Initializers.Initializer):

+class Bernouilli(tf.keras.initializers.Initializer):+  """+  Initialization distributio following a Bernouilli process..+  """++  def __init__(self, p):+    """+    p: probability parameter of success (i.e. 1).+    """+    self.p = p++  def get_config(self):+    return {'p': self.p}++  @classmethod+  def from_config(cls, config):+    return cls(**config)++  def __call__(self, shape, dytpe=tf.dtypes.float32):+    """Number of zeros = np.ceil(sparsity * size) in expectation."""+    probs = tf.zeros(shape=list(shape)) + self.p+    uniform = tf.random.uniform(shape)+    initial = tf.less(uniform, probs)++    return tf.cast(initial, dtype=dtype)++class PermuteOnes(tf.keras.initializers.Initializer):+  """+  Initialization of a deterministically sparse matrix.+  """+  def __init__(self, p=None):+    """+    p: probability parameter of success (i.e. 1).

+class Bernouilli(tf.keras.initializers.Initializer):+  """+  Initialization distributio following a Bernouilli process..+  """++  def __init__(self, p):+    """+    p: probability parameter of success (i.e. 1).+    """+    self.p = p++  def get_config(self):+    return {'p': self.p}++  @classmethod+  def from_config(cls, config):+    return cls(**config)++  def __call__(self, shape, dytpe=tf.dtypes.float32):+    """Number of zeros = np.ceil(sparsity * size) in expectation."""+    probs = tf.zeros(shape=list(shape)) + self.p+    uniform = tf.random.uniform(shape)+    initial = tf.less(uniform, probs)++    return tf.cast(initial, dtype=dtype)++class PermuteOnes(tf.keras.initializers.Initializer):+  """+  Initialization of a deterministically sparse matrix.

+class Bernouilli(tf.keras.initializers.Initializer):+  """+  Initialization distributio following a Bernouilli process..+  """++  def __init__(self, p):+    """+    p: probability parameter of success (i.e. 1).+    """+    self.p = p++  def get_config(self):+    return {'p': self.p}++  @classmethod+  def from_config(cls, config):+    return cls(**config)++  def __call__(self, shape, dytpe=tf.dtypes.float32):+    """Number of zeros = np.ceil(sparsity * size) in expectation."""+    probs = tf.zeros(shape=list(shape)) + self.p+    uniform = tf.random.uniform(shape)+    initial = tf.less(uniform, probs)++    return tf.cast(initial, dtype=dtype)++class PermuteOnes(tf.keras.initializers.Initializer):+  """+  Initialization of a deterministically sparse matrix.+  """+  def __init__(self, p=None):+    """+    p: probability parameter of success (i.e. 1).+    If p is None, will sample randomly from uniform distribution for sparsity.+    """+    self.p = p if p else tf.random.uniform(())++  def get_n_ones(self, shape, dtype=tf.dtypes.float32):+    sparsity = self.p if self.p else 0.0+    return tf.math.ceil(sparsity * tf.cast(tf.math.reduce_sum(shape)), dtype=dtype)++  def __call__(self, shape, dtype=tf.dtypes.float32, seed=None):+    flat_mask = tf.reshape(tf.ones(shape), (-1,))+    num_ones = self.get_n_ones(shape, dtype)+    _indices = tf.cast(tf.reshape(tf.linspace(0, num_ones - 1, int(num_ones)), (-1,)), tf.int32)+    indices = tf.reshape(_indices, (-1, 1))+    updates = tf.ones_like(_indices)+    flat_shape = flat_mask.shape+    unshuffled_mask = tf.scatter_nd(indices, udpates, flat_shape)+    shuffled_mask = tf.random.shuffle(unshuffled_mas, seed=seed)

++# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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.+# ==============================================================================+"""Helper functions to add support for iterative magnitude/gradient pruning as seen in the RiGL experiments."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++import tensorflow as tf+from absl import logging++from tensorflow.python.ops import summary_ops_v2+from tensorflow.python.summary import summary as summary_ops_v1+from tensorflow_model_optimization.python.core.sparsity.keras import pruning_utils+from tensorflow_model_optimization.python.core.sparsity_tf2 import schedule as update_schedule+from tensorflow_model_optimization.python.core.sparsity_tf2 import pruner+from tensorflow_model_optimization.python.core.sparsity_tf2 import sparse_utils+++class RiGLPruner(pruner.Pruner):+  """+  Implementation of the RiGL dynamic sparse training algorithm.+  """++  def __init__(self,+      update_schedule=update_schedule.ConstantSchedule(0.5, 0),+      sparsity=0.5,+      block_size=(1,1),+      block_pooling_type='AVG',+      initializer=sparse_utils.PermuteOnes,+      stateless=False,+      seed=0,+      seed_offset=0,+      noise_std=0,+      reinit=False+    ):+    """The logic for magnitude-based pruning weight tensors.++    Args:+      update_schedule: A `PruningSchedule` object that controls pruning rate+        throughout training.+      sparsity: the sparsity at which the dynamic sparse training method uses+      block_size: The dimensions (height, weight) for the block sparse pattern+        in rank-2 weight tensors.+      block_pooling_type: (optional) The function to use to pool weights in the+        block. Must be 'AVG' or 'MAX'.+      initializer: the initial sparsity distribution Callable for each layer of the network.+      stateless: whether or not being run on TPU/multi-GPU workers.+      seed: assigned by PruningConfig for multiworker consistency in random processes.+      seed_offset: added to seed to run different experiments.+      reinit: boolean for whether to reinitialize a connection that was drop and regorwn to +        its original value, or to set it to 0+    """+    super(RiGLPruner, self).__init__(+        update_schedule, block_size, block_pooling_type)+    self.initializer = initializer+    self.target_sparsity = sparsity+    self.update_schedule = update_schedule+    self._block_size = block_size+    self.block_pooling_type = block_pooling_type+    self._stateless = stateless+    self._seed = seed+    self._seed_offset = seed_offset+    self._noise_std = noise_std+    self._reinit_when_same = reinit+  +  def create_slots(self, optimizer, var):+    base_dtype = var.dtype+    optimizer.add_slot(var, 'mask', initializer=self.initializer(self.target_sparsity))+++  def _validate_block(self, var):+    if self._block_size != [1, 1]:+      if var.get_shape().ndims != 2:+        raise ValueError('Block Sparsity can only be used for layers which '+                          'have 2-dimensional weights.')++  def _random_uniform(self, step, *args, **kwargs):+    """Uniform noise distribution"""+    if self._stateless:+      offset_seed = self._seed_offset + kwargs.get('seed', 0)+      kwargs['seed'] = tf.cast(+        tf.stack([offset_seed, step], tf.int32)+      )+      return tf.random.stateless_uniform(*args, **kwargs)+    return tf.random.uniform(*args, **kwargs)++  def _random_normal(self, step, *args, **kwargs):+    """Gaussian noise distribution"""+    if self._stateless:+      offset_seed = self._seed_offset + kwargs.get('seed', 0)+      kwargs['seed'] = tf.cast(+        tf.stack([offset_seed, step], tf.int32)+      )+      return tf.random.stateless_normal(*args, **kwargs)+    return tf.random.normal(*args, **kwargs)++  def _get_grow_grads(self, optimizer, step, mask, grad):+    """Grow connections baased on gradient information.+    """+    grad_scores = tf.math.abs(grad)+    return grad_scores++  def _get_drop_weights(self, optimizer, step, mask, weight, noise_std=0):+    """+    Drop connections based on weight magnitude.+    """+    masked_weight = mask * weight+    weight_scores = tf.math.abs(masked_weight)+    # add noise to break ties, although the possibility is very low+    if noise_std != 0:+      weight_scores.assign_add(self._random_normal(+        weight_scores.shape, stddev=noise_std, dtype=weight_scores.dtype,+        seed=(hash(weight.name + 'drop'))))+    return weight_scores++  def _reset_momentum(self, optimizer, weight, new_connections):+    """Zeros out optimizer slots whose connections have been recovered."""+    for slot_name in optimizer._optimizer.get_slot_names():+      # reset aggregated momentum variables to 0+      opt_var = optimizer._optimizer.get_slot(weight, slot_name)+      new_values = tf.where(new_connections,+                            tf.zeros_like(opt_var), opt_var)+      opt_var.assign(new_values)++  def _generic_top_k(self, scores, mask, n_new, n_total):+     # sort the entire array so that it can be constant for TPUs+    _, sorted_idx = tf.math.top_k(+      tf.reshape(scores, (-1), k=n_total)+    )+    expanded_sorted_idx = tf.expand_dims(sorted_idx, 1)+    new_values = tf.where(+      tf.range(n_total) < n_new,+      tf.ones_like(sorted_idx, dtype=mask.dtype),+      tf.zeros_like(sorted_idx, dtype=mask.dtype)+    )+    updated_mask = tf.scatter_nd(expanded_sorted_idx, new_values, new_values.shape)++    return updated_mask++  def _get_new_connections(self, reinit_when_same, grown_mask_reshaped, mask):+    if reinit_when_same:+      new_connections = tf.math.equal(grown_mask_reshaped, 1)+    else:+      new_connections = tf.math.logical_and(+        tf.math.equal(grown_mask_reshaped, 1), tf.math.equal(mask, 0)+      )+    return new_connections++  def _update_mask(self, step, update_fraction, mask, weight, grad):+    """Called by _maybe_update_block_mask.++    Updates mask based on weight and grad information.+    """+    # compute the top k magnitudes then update the current mask+    drop_scores = self._get_drop_weights(optimizer, step, mask, weight, noise_std=self._noise_std)+    # need access to exactly which entries are growing to zero out optimizer slot+    grow_scores = self._get_grow_grads(optimizer, step, mask, grad)+    n_total = tf.size(drop_scores)+    n_ones = tf.cast(tf.reduce_sum(mask), dtype=tf.int32)

++# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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.+# ==============================================================================+"""Helper functions to add support for iterative magnitude/gradient pruning as seen in the RiGL experiments."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++import tensorflow as tf+from absl import logging++from tensorflow.python.ops import summary_ops_v2+from tensorflow.python.summary import summary as summary_ops_v1+from tensorflow_model_optimization.python.core.sparsity.keras import pruning_utils+from tensorflow_model_optimization.python.core.sparsity_tf2 import schedule as update_schedule+from tensorflow_model_optimization.python.core.sparsity_tf2 import pruner+from tensorflow_model_optimization.python.core.sparsity_tf2 import sparse_utils+++class RiGLPruner(pruner.Pruner):+  """+  Implementation of the RiGL dynamic sparse training algorithm.+  """++  def __init__(self,+      update_schedule=update_schedule.ConstantSchedule(0.5, 0),+      sparsity=0.5,+      block_size=(1,1),+      block_pooling_type='AVG',+      initializer=sparse_utils.PermuteOnes,+      stateless=False,+      seed=0,+      seed_offset=0,+      noise_std=0,+      reinit=False+    ):+    """The logic for magnitude-based pruning weight tensors.++    Args:+      update_schedule: A `PruningSchedule` object that controls pruning rate+        throughout training.+      sparsity: the sparsity at which the dynamic sparse training method uses+      block_size: The dimensions (height, weight) for the block sparse pattern+        in rank-2 weight tensors.+      block_pooling_type: (optional) The function to use to pool weights in the+        block. Must be 'AVG' or 'MAX'.+      initializer: the initial sparsity distribution Callable for each layer of the network.+      stateless: whether or not being run on TPU/multi-GPU workers.+      seed: assigned by PruningConfig for multiworker consistency in random processes.+      seed_offset: added to seed to run different experiments.+      reinit: boolean for whether to reinitialize a connection that was drop and regorwn to +        its original value, or to set it to 0+    """+    super(RiGLPruner, self).__init__(+        update_schedule, block_size, block_pooling_type)+    self.initializer = initializer+    self.target_sparsity = sparsity+    self.update_schedule = update_schedule+    self._block_size = block_size+    self.block_pooling_type = block_pooling_type+    self._stateless = stateless+    self._seed = seed+    self._seed_offset = seed_offset+    self._noise_std = noise_std+    self._reinit_when_same = reinit+  +  def create_slots(self, optimizer, var):+    base_dtype = var.dtype+    optimizer.add_slot(var, 'mask', initializer=self.initializer(self.target_sparsity))+++  def _validate_block(self, var):+    if self._block_size != [1, 1]:+      if var.get_shape().ndims != 2:+        raise ValueError('Block Sparsity can only be used for layers which '+                          'have 2-dimensional weights.')++  def _random_uniform(self, step, *args, **kwargs):+    """Uniform noise distribution"""+    if self._stateless:+      offset_seed = self._seed_offset + kwargs.get('seed', 0)+      kwargs['seed'] = tf.cast(+        tf.stack([offset_seed, step], tf.int32)+      )+      return tf.random.stateless_uniform(*args, **kwargs)+    return tf.random.uniform(*args, **kwargs)++  def _random_normal(self, step, *args, **kwargs):+    """Gaussian noise distribution"""+    if self._stateless:+      offset_seed = self._seed_offset + kwargs.get('seed', 0)+      kwargs['seed'] = tf.cast(+        tf.stack([offset_seed, step], tf.int32)+      )+      return tf.random.stateless_normal(*args, **kwargs)+    return tf.random.normal(*args, **kwargs)++  def _get_grow_grads(self, optimizer, step, mask, grad):+    """Grow connections baased on gradient information.+    """+    grad_scores = tf.math.abs(grad)+    return grad_scores++  def _get_drop_weights(self, optimizer, step, mask, weight, noise_std=0):+    """+    Drop connections based on weight magnitude.+    """+    masked_weight = mask * weight+    weight_scores = tf.math.abs(masked_weight)+    # add noise to break ties, although the possibility is very low+    if noise_std != 0:+      weight_scores.assign_add(self._random_normal(+        weight_scores.shape, stddev=noise_std, dtype=weight_scores.dtype,+        seed=(hash(weight.name + 'drop'))))+    return weight_scores++  def _reset_momentum(self, optimizer, weight, new_connections):+    """Zeros out optimizer slots whose connections have been recovered."""+    for slot_name in optimizer._optimizer.get_slot_names():+      # reset aggregated momentum variables to 0+      opt_var = optimizer._optimizer.get_slot(weight, slot_name)+      new_values = tf.where(new_connections,+                            tf.zeros_like(opt_var), opt_var)+      opt_var.assign(new_values)++  def _generic_top_k(self, scores, mask, n_new, n_total):+     # sort the entire array so that it can be constant for TPUs+    _, sorted_idx = tf.math.top_k(+      tf.reshape(scores, (-1), k=n_total)+    )+    expanded_sorted_idx = tf.expand_dims(sorted_idx, 1)+    new_values = tf.where(+      tf.range(n_total) < n_new,+      tf.ones_like(sorted_idx, dtype=mask.dtype),+      tf.zeros_like(sorted_idx, dtype=mask.dtype)+    )+    updated_mask = tf.scatter_nd(expanded_sorted_idx, new_values, new_values.shape)++    return updated_mask++  def _get_new_connections(self, reinit_when_same, grown_mask_reshaped, mask):+    if reinit_when_same:+      new_connections = tf.math.equal(grown_mask_reshaped, 1)+    else:+      new_connections = tf.math.logical_and(+        tf.math.equal(grown_mask_reshaped, 1), tf.math.equal(mask, 0)+      )+    return new_connections++  def _update_mask(self, step, update_fraction, mask, weight, grad):+    """Called by _maybe_update_block_mask.++    Updates mask based on weight and grad information.+    """+    # compute the top k magnitudes then update the current mask+    drop_scores = self._get_drop_weights(optimizer, step, mask, weight, noise_std=self._noise_std)+    # need access to exactly which entries are growing to zero out optimizer slot+    grow_scores = self._get_grow_grads(optimizer, step, mask, grad)+    n_total = tf.size(drop_scores)+    n_ones = tf.cast(tf.reduce_sum(mask), dtype=tf.int32)+    n_prune = tf.cast(+      tf.cast(n_ones, dtype=tf.float32) * update_fraction, tf.int32+    )+    n_keep = n_ones - n_prune++    dropped_mask = self._generic_top_k(drop_scores, mask, n_keep, n_total)++    if grow_scores is not None:+      # flatten the scores+      grow_scores = tf.reshape(grow_scores, (-1))+      # set enabled connections (ones) to min(scores) - 1, i.e. they have the lowest scores+      grow_scores_lifted = tf.where(+        tf.math.equal(dropped_mask, 1),+        tf.ones_like(dropped_mask) * (tf.reduce_min(grow_scores) - 1), grow_scores+      )+      grown_mask = self._generic_top_k(grow_scores_lifted, mask, n_prune, n_total)+      # ensure that masks are disjoint

++# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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.+# ==============================================================================+"""Helper functions to add support for iterative magnitude/gradient pruning as seen in the RiGL experiments."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++import tensorflow as tf+from absl import logging++from tensorflow.python.ops import summary_ops_v2+from tensorflow.python.summary import summary as summary_ops_v1+from tensorflow_model_optimization.python.core.sparsity.keras import pruning_utils+from tensorflow_model_optimization.python.core.sparsity_tf2 import schedule as update_schedule+from tensorflow_model_optimization.python.core.sparsity_tf2 import pruner+from tensorflow_model_optimization.python.core.sparsity_tf2 import sparse_utils+++class RiGLPruner(pruner.Pruner):+  """+  Implementation of the RiGL dynamic sparse training algorithm.+  """++  def __init__(self,+      update_schedule=update_schedule.ConstantSchedule(0.5, 0),+      sparsity=0.5,+      block_size=(1,1),+      block_pooling_type='AVG',+      initializer=sparse_utils.PermuteOnes,+      stateless=False,+      seed=0,+      seed_offset=0,+      noise_std=0,+      reinit=False+    ):+    """The logic for magnitude-based pruning weight tensors.

++# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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.+# ==============================================================================+"""Helper functions to add support for iterative magnitude/gradient pruning as seen in the RiGL experiments."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++import tensorflow as tf+from absl import logging++from tensorflow.python.ops import summary_ops_v2+from tensorflow.python.summary import summary as summary_ops_v1+from tensorflow_model_optimization.python.core.sparsity.keras import pruning_utils+from tensorflow_model_optimization.python.core.sparsity_tf2 import schedule as update_schedule+from tensorflow_model_optimization.python.core.sparsity_tf2 import pruner+from tensorflow_model_optimization.python.core.sparsity_tf2 import sparse_utils+++class RiGLPruner(pruner.Pruner):+  """+  Implementation of the RiGL dynamic sparse training algorithm.+  """++  def __init__(self,+      update_schedule=update_schedule.ConstantSchedule(0.5, 0),+      sparsity=0.5,+      block_size=(1,1),+      block_pooling_type='AVG',+      initializer=sparse_utils.PermuteOnes,+      stateless=False,+      seed=0,+      seed_offset=0,+      noise_std=0,+      reinit=False+    ):+    """The logic for magnitude-based pruning weight tensors.++    Args:+      update_schedule: A `PruningSchedule` object that controls pruning rate+        throughout training.+      sparsity: the sparsity at which the dynamic sparse training method uses+      block_size: The dimensions (height, weight) for the block sparse pattern+        in rank-2 weight tensors.+      block_pooling_type: (optional) The function to use to pool weights in the+        block. Must be 'AVG' or 'MAX'.+      initializer: the initial sparsity distribution Callable for each layer of the network.+      stateless: whether or not being run on TPU/multi-GPU workers.+      seed: assigned by PruningConfig for multiworker consistency in random processes.+      seed_offset: added to seed to run different experiments.+      reinit: boolean for whether to reinitialize a connection that was drop and regorwn to +        its original value, or to set it to 0+    """+    super(RiGLPruner, self).__init__(+        update_schedule, block_size, block_pooling_type)+    self.initializer = initializer+    self.target_sparsity = sparsity+    self.update_schedule = update_schedule+    self._block_size = block_size+    self.block_pooling_type = block_pooling_type+    self._stateless = stateless+    self._seed = seed+    self._seed_offset = seed_offset+    self._noise_std = noise_std+    self._reinit_when_same = reinit+  +  def create_slots(self, optimizer, var):+    base_dtype = var.dtype+    optimizer.add_slot(var, 'mask', initializer=self.initializer(self.target_sparsity))+++  def _validate_block(self, var):+    if self._block_size != [1, 1]:+      if var.get_shape().ndims != 2:+        raise ValueError('Block Sparsity can only be used for layers which '+                          'have 2-dimensional weights.')++  def _random_uniform(self, step, *args, **kwargs):+    """Uniform noise distribution"""+    if self._stateless:+      offset_seed = self._seed_offset + kwargs.get('seed', 0)+      kwargs['seed'] = tf.cast(+        tf.stack([offset_seed, step], tf.int32)+      )+      return tf.random.stateless_uniform(*args, **kwargs)+    return tf.random.uniform(*args, **kwargs)++  def _random_normal(self, step, *args, **kwargs):+    """Gaussian noise distribution"""+    if self._stateless:+      offset_seed = self._seed_offset + kwargs.get('seed', 0)+      kwargs['seed'] = tf.cast(+        tf.stack([offset_seed, step], tf.int32)+      )+      return tf.random.stateless_normal(*args, **kwargs)+    return tf.random.normal(*args, **kwargs)++  def _get_grow_grads(self, optimizer, step, mask, grad):+    """Grow connections baased on gradient information.+    """+    grad_scores = tf.math.abs(grad)+    return grad_scores++  def _get_drop_weights(self, optimizer, step, mask, weight, noise_std=0):+    """+    Drop connections based on weight magnitude.+    """+    masked_weight = mask * weight+    weight_scores = tf.math.abs(masked_weight)+    # add noise to break ties, although the possibility is very low+    if noise_std != 0:+      weight_scores.assign_add(self._random_normal(+        weight_scores.shape, stddev=noise_std, dtype=weight_scores.dtype,+        seed=(hash(weight.name + 'drop'))))+    return weight_scores++  def _reset_momentum(self, optimizer, weight, new_connections):+    """Zeros out optimizer slots whose connections have been recovered."""

+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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 key functions in riglpruner library."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++# import g3++from absl.testing import parameterized+import numpy as np+import tensorflow as tf++# TODO(b/139939526): move to public API.+from tensorflow.python.keras import keras_parameterized+from tensorflow_model_optimization.python.core.sparsity.keras import pruning_schedule+from tensorflow_model_optimization.python.core.sparsity_tf2 import riglpruner as pruner++dtypes = tf.dtypes+test = tf.test+++def make_pruning_schedule(target_sparsity, begin, end, freq):+  return pruning_schedule.ConstantSparsity(target_sparsity, begin, end, freq)++def sample_noise(x, mu=0, sigma=1.):+  sample = tf.random.normal((), mean=mu,  stddev=sigma, dtype=tf.float64)+  return sample++def _dummy_gradient(x, dtype=tf.float32):+  try:+    base_type = x.dtype+  except:+    base_type = dtype+  grad = tf.ones_like(x, dtype=base_type)+  return grad++class RiglPruningTest(test.TestCase, parameterized.TestCase):

++# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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.+# ==============================================================================+"""Helper functions to add support for iterative magnitude/gradient pruning as seen in the RiGL experiments."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++import tensorflow as tf+from absl import logging++from tensorflow.python.ops import summary_ops_v2+from tensorflow.python.summary import summary as summary_ops_v1+from tensorflow_model_optimization.python.core.sparsity.keras import pruning_utils+from tensorflow_model_optimization.python.core.sparsity_tf2 import schedule as update_schedule+from tensorflow_model_optimization.python.core.sparsity_tf2 import pruner+from tensorflow_model_optimization.python.core.sparsity_tf2 import sparse_utils+++class RiGLPruner(pruner.Pruner):+  """+  Implementation of the RiGL dynamic sparse training algorithm.+  """++  def __init__(self,+      update_schedule=update_schedule.ConstantSchedule(0.5, 0),+      sparsity=0.5,+      block_size=(1,1),+      block_pooling_type='AVG',+      initializer=sparse_utils.PermuteOnes,+      stateless=False,+      seed=0,+      seed_offset=0,+      noise_std=0,+      reinit=False+    ):+    """The logic for magnitude-based pruning weight tensors.++    Args:+      update_schedule: A `PruningSchedule` object that controls pruning rate+        throughout training.+      sparsity: the sparsity at which the dynamic sparse training method uses+      block_size: The dimensions (height, weight) for the block sparse pattern+        in rank-2 weight tensors.+      block_pooling_type: (optional) The function to use to pool weights in the+        block. Must be 'AVG' or 'MAX'.+      initializer: the initial sparsity distribution Callable for each layer of the network.+      stateless: whether or not being run on TPU/multi-GPU workers.+      seed: assigned by PruningConfig for multiworker consistency in random processes.+      seed_offset: added to seed to run different experiments.+      reinit: boolean for whether to reinitialize a connection that was drop and regorwn to +        its original value, or to set it to 0+    """+    super(RiGLPruner, self).__init__(+        update_schedule, block_size, block_pooling_type)+    self.initializer = initializer+    self.target_sparsity = sparsity+    self.update_schedule = update_schedule+    self._block_size = block_size+    self.block_pooling_type = block_pooling_type+    self._stateless = stateless+    self._seed = seed+    self._seed_offset = seed_offset+    self._noise_std = noise_std+    self._reinit_when_same = reinit+  +  def create_slots(self, optimizer, var):+    base_dtype = var.dtype+    optimizer.add_slot(var, 'mask', initializer=self.initializer(self.target_sparsity))+++  def _validate_block(self, var):+    if self._block_size != [1, 1]:+      if var.get_shape().ndims != 2:+        raise ValueError('Block Sparsity can only be used for layers which '+                          'have 2-dimensional weights.')++  def _random_uniform(self, step, *args, **kwargs):

++# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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.+# ==============================================================================+"""Helper functions to add support for iterative magnitude/gradient pruning as seen in the RiGL experiments."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++import tensorflow as tf+from absl import logging++from tensorflow.python.ops import summary_ops_v2+from tensorflow.python.summary import summary as summary_ops_v1+from tensorflow_model_optimization.python.core.sparsity.keras import pruning_utils+from tensorflow_model_optimization.python.core.sparsity_tf2 import schedule as update_schedule+from tensorflow_model_optimization.python.core.sparsity_tf2 import pruner+from tensorflow_model_optimization.python.core.sparsity_tf2 import sparse_utils+++class RiGLPruner(pruner.Pruner):+  """+  Implementation of the RiGL dynamic sparse training algorithm.+  """++  def __init__(self,+      update_schedule=update_schedule.ConstantSchedule(0.5, 0),+      sparsity=0.5,+      block_size=(1,1),+      block_pooling_type='AVG',+      initializer=sparse_utils.PermuteOnes,+      stateless=False,+      seed=0,+      seed_offset=0,+      noise_std=0,+      reinit=False+    ):+    """The logic for magnitude-based pruning weight tensors.++    Args:+      update_schedule: A `PruningSchedule` object that controls pruning rate+        throughout training.+      sparsity: the sparsity at which the dynamic sparse training method uses+      block_size: The dimensions (height, weight) for the block sparse pattern+        in rank-2 weight tensors.+      block_pooling_type: (optional) The function to use to pool weights in the+        block. Must be 'AVG' or 'MAX'.+      initializer: the initial sparsity distribution Callable for each layer of the network.+      stateless: whether or not being run on TPU/multi-GPU workers.+      seed: assigned by PruningConfig for multiworker consistency in random processes.+      seed_offset: added to seed to run different experiments.+      reinit: boolean for whether to reinitialize a connection that was drop and regorwn to +        its original value, or to set it to 0+    """+    super(RiGLPruner, self).__init__(+        update_schedule, block_size, block_pooling_type)+    self.initializer = initializer+    self.target_sparsity = sparsity+    self.update_schedule = update_schedule+    self._block_size = block_size+    self.block_pooling_type = block_pooling_type+    self._stateless = stateless+    self._seed = seed+    self._seed_offset = seed_offset+    self._noise_std = noise_std+    self._reinit_when_same = reinit+  +  def create_slots(self, optimizer, var):+    base_dtype = var.dtype+    optimizer.add_slot(var, 'mask', initializer=self.initializer(self.target_sparsity))+++  def _validate_block(self, var):+    if self._block_size != [1, 1]:+      if var.get_shape().ndims != 2:+        raise ValueError('Block Sparsity can only be used for layers which '+                          'have 2-dimensional weights.')++  def _random_uniform(self, step, *args, **kwargs):+    """Uniform noise distribution"""+    if self._stateless:+      offset_seed = self._seed_offset + kwargs.get('seed', 0)+      kwargs['seed'] = tf.cast(+        tf.stack([offset_seed, step], tf.int32)+      )+      return tf.random.stateless_uniform(*args, **kwargs)+    return tf.random.uniform(*args, **kwargs)++  def _random_normal(self, step, *args, **kwargs):+    """Gaussian noise distribution"""+    if self._stateless:+      offset_seed = self._seed_offset + kwargs.get('seed', 0)

++# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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.+# ==============================================================================+"""Helper functions to add support for iterative magnitude/gradient pruning as seen in the RiGL experiments."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++import tensorflow as tf+from absl import logging++from tensorflow.python.ops import summary_ops_v2+from tensorflow.python.summary import summary as summary_ops_v1+from tensorflow_model_optimization.python.core.sparsity.keras import pruning_utils+from tensorflow_model_optimization.python.core.sparsity_tf2 import schedule as update_schedule+from tensorflow_model_optimization.python.core.sparsity_tf2 import pruner+from tensorflow_model_optimization.python.core.sparsity_tf2 import sparse_utils+++class RiGLPruner(pruner.Pruner):+  """+  Implementation of the RiGL dynamic sparse training algorithm.+  """++  def __init__(self,+      update_schedule=update_schedule.ConstantSchedule(0.5, 0),+      sparsity=0.5,+      block_size=(1,1),+      block_pooling_type='AVG',+      initializer=sparse_utils.PermuteOnes,+      stateless=False,+      seed=0,+      seed_offset=0,+      noise_std=0,+      reinit=False+    ):+    """The logic for magnitude-based pruning weight tensors.++    Args:+      update_schedule: A `PruningSchedule` object that controls pruning rate+        throughout training.+      sparsity: the sparsity at which the dynamic sparse training method uses+      block_size: The dimensions (height, weight) for the block sparse pattern+        in rank-2 weight tensors.+      block_pooling_type: (optional) The function to use to pool weights in the+        block. Must be 'AVG' or 'MAX'.+      initializer: the initial sparsity distribution Callable for each layer of the network.+      stateless: whether or not being run on TPU/multi-GPU workers.+      seed: assigned by PruningConfig for multiworker consistency in random processes.+      seed_offset: added to seed to run different experiments.+      reinit: boolean for whether to reinitialize a connection that was drop and regorwn to +        its original value, or to set it to 0+    """+    super(RiGLPruner, self).__init__(+        update_schedule, block_size, block_pooling_type)+    self.initializer = initializer+    self.target_sparsity = sparsity+    self.update_schedule = update_schedule+    self._block_size = block_size+    self.block_pooling_type = block_pooling_type+    self._stateless = stateless+    self._seed = seed+    self._seed_offset = seed_offset+    self._noise_std = noise_std+    self._reinit_when_same = reinit+  +  def create_slots(self, optimizer, var):+    base_dtype = var.dtype+    optimizer.add_slot(var, 'mask', initializer=self.initializer(self.target_sparsity))+++  def _validate_block(self, var):+    if self._block_size != [1, 1]:+      if var.get_shape().ndims != 2:+        raise ValueError('Block Sparsity can only be used for layers which '+                          'have 2-dimensional weights.')++  def _random_uniform(self, step, *args, **kwargs):+    """Uniform noise distribution"""+    if self._stateless:+      offset_seed = self._seed_offset + kwargs.get('seed', 0)+      kwargs['seed'] = tf.cast(+        tf.stack([offset_seed, step], tf.int32)+      )+      return tf.random.stateless_uniform(*args, **kwargs)+    return tf.random.uniform(*args, **kwargs)++  def _random_normal(self, step, *args, **kwargs):+    """Gaussian noise distribution"""+    if self._stateless:+      offset_seed = self._seed_offset + kwargs.get('seed', 0)+      kwargs['seed'] = tf.cast(+        tf.stack([offset_seed, step], tf.int32)+      )+      return tf.random.stateless_normal(*args, **kwargs)+    return tf.random.normal(*args, **kwargs)++  def _get_grow_grads(self, optimizer, step, mask, grad):+    """Grow connections baased on gradient information.+    """+    grad_scores = tf.math.abs(grad)+    return grad_scores++  def _get_drop_weights(self, optimizer, step, mask, weight, noise_std=0):+    """+    Drop connections based on weight magnitude.+    """+    masked_weight = mask * weight+    weight_scores = tf.math.abs(masked_weight)+    # add noise to break ties, although the possibility is very low+    if noise_std != 0:+      weight_scores.assign_add(self._random_normal(+        weight_scores.shape, stddev=noise_std, dtype=weight_scores.dtype,+        seed=(hash(weight.name + 'drop'))))+    return weight_scores++  def _reset_momentum(self, optimizer, weight, new_connections):+    """Zeros out optimizer slots whose connections have been recovered."""+    for slot_name in optimizer._optimizer.get_slot_names():+      # reset aggregated momentum variables to 0+      opt_var = optimizer._optimizer.get_slot(weight, slot_name)+      new_values = tf.where(new_connections,

++# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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.+# ==============================================================================+"""Helper functions to add support for iterative magnitude/gradient pruning as seen in the RiGL experiments."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++import tensorflow as tf+from absl import logging++from tensorflow.python.ops import summary_ops_v2+from tensorflow.python.summary import summary as summary_ops_v1+from tensorflow_model_optimization.python.core.sparsity.keras import pruning_utils+from tensorflow_model_optimization.python.core.sparsity_tf2 import schedule as update_schedule+from tensorflow_model_optimization.python.core.sparsity_tf2 import pruner+from tensorflow_model_optimization.python.core.sparsity_tf2 import sparse_utils+++class RiGLPruner(pruner.Pruner):+  """+  Implementation of the RiGL dynamic sparse training algorithm.+  """++  def __init__(self,+      update_schedule=update_schedule.ConstantSchedule(0.5, 0),+      sparsity=0.5,+      block_size=(1,1),+      block_pooling_type='AVG',+      initializer=sparse_utils.PermuteOnes,+      stateless=False,+      seed=0,+      seed_offset=0,+      noise_std=0,+      reinit=False+    ):+    """The logic for magnitude-based pruning weight tensors.++    Args:+      update_schedule: A `PruningSchedule` object that controls pruning rate+        throughout training.+      sparsity: the sparsity at which the dynamic sparse training method uses+      block_size: The dimensions (height, weight) for the block sparse pattern+        in rank-2 weight tensors.+      block_pooling_type: (optional) The function to use to pool weights in the+        block. Must be 'AVG' or 'MAX'.+      initializer: the initial sparsity distribution Callable for each layer of the network.+      stateless: whether or not being run on TPU/multi-GPU workers.+      seed: assigned by PruningConfig for multiworker consistency in random processes.+      seed_offset: added to seed to run different experiments.+      reinit: boolean for whether to reinitialize a connection that was drop and regorwn to +        its original value, or to set it to 0+    """+    super(RiGLPruner, self).__init__(+        update_schedule, block_size, block_pooling_type)+    self.initializer = initializer+    self.target_sparsity = sparsity+    self.update_schedule = update_schedule+    self._block_size = block_size+    self.block_pooling_type = block_pooling_type+    self._stateless = stateless+    self._seed = seed+    self._seed_offset = seed_offset+    self._noise_std = noise_std+    self._reinit_when_same = reinit+  +  def create_slots(self, optimizer, var):+    base_dtype = var.dtype+    optimizer.add_slot(var, 'mask', initializer=self.initializer(self.target_sparsity))+++  def _validate_block(self, var):+    if self._block_size != [1, 1]:+      if var.get_shape().ndims != 2:+        raise ValueError('Block Sparsity can only be used for layers which '+                          'have 2-dimensional weights.')++  def _random_uniform(self, step, *args, **kwargs):+    """Uniform noise distribution"""+    if self._stateless:+      offset_seed = self._seed_offset + kwargs.get('seed', 0)+      kwargs['seed'] = tf.cast(+        tf.stack([offset_seed, step], tf.int32)+      )+      return tf.random.stateless_uniform(*args, **kwargs)+    return tf.random.uniform(*args, **kwargs)++  def _random_normal(self, step, *args, **kwargs):+    """Gaussian noise distribution"""+    if self._stateless:+      offset_seed = self._seed_offset + kwargs.get('seed', 0)+      kwargs['seed'] = tf.cast(+        tf.stack([offset_seed, step], tf.int32)+      )+      return tf.random.stateless_normal(*args, **kwargs)+    return tf.random.normal(*args, **kwargs)++  def _get_grow_grads(self, optimizer, step, mask, grad):+    """Grow connections baased on gradient information.+    """+    grad_scores = tf.math.abs(grad)+    return grad_scores++  def _get_drop_weights(self, optimizer, step, mask, weight, noise_std=0):+    """+    Drop connections based on weight magnitude.+    """+    masked_weight = mask * weight+    weight_scores = tf.math.abs(masked_weight)+    # add noise to break ties, although the possibility is very low+    if noise_std != 0:+      weight_scores.assign_add(self._random_normal(+        weight_scores.shape, stddev=noise_std, dtype=weight_scores.dtype,+        seed=(hash(weight.name + 'drop'))))

++# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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.+# ==============================================================================+"""Helper functions to add support for iterative magnitude/gradient pruning as seen in the RiGL experiments."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++import tensorflow as tf+from absl import logging++from tensorflow.python.ops import summary_ops_v2+from tensorflow.python.summary import summary as summary_ops_v1+from tensorflow_model_optimization.python.core.sparsity.keras import pruning_utils+from tensorflow_model_optimization.python.core.sparsity_tf2 import schedule as update_schedule+from tensorflow_model_optimization.python.core.sparsity_tf2 import pruner+from tensorflow_model_optimization.python.core.sparsity_tf2 import sparse_utils+++class RiGLPruner(pruner.Pruner):+  """+  Implementation of the RiGL dynamic sparse training algorithm.+  """++  def __init__(self,+      update_schedule=update_schedule.ConstantSchedule(0.5, 0),+      sparsity=0.5,+      block_size=(1,1),+      block_pooling_type='AVG',+      initializer=sparse_utils.PermuteOnes,+      stateless=False,+      seed=0,+      seed_offset=0,+      noise_std=0,+      reinit=False+    ):+    """The logic for magnitude-based pruning weight tensors.++    Args:+      update_schedule: A `PruningSchedule` object that controls pruning rate+        throughout training.+      sparsity: the sparsity at which the dynamic sparse training method uses+      block_size: The dimensions (height, weight) for the block sparse pattern+        in rank-2 weight tensors.+      block_pooling_type: (optional) The function to use to pool weights in the+        block. Must be 'AVG' or 'MAX'.+      initializer: the initial sparsity distribution Callable for each layer of the network.+      stateless: whether or not being run on TPU/multi-GPU workers.+      seed: assigned by PruningConfig for multiworker consistency in random processes.+      seed_offset: added to seed to run different experiments.+      reinit: boolean for whether to reinitialize a connection that was drop and regorwn to +        its original value, or to set it to 0+    """+    super(RiGLPruner, self).__init__(+        update_schedule, block_size, block_pooling_type)+    self.initializer = initializer+    self.target_sparsity = sparsity+    self.update_schedule = update_schedule+    self._block_size = block_size+    self.block_pooling_type = block_pooling_type+    self._stateless = stateless+    self._seed = seed+    self._seed_offset = seed_offset+    self._noise_std = noise_std+    self._reinit_when_same = reinit+  +  def create_slots(self, optimizer, var):+    base_dtype = var.dtype+    optimizer.add_slot(var, 'mask', initializer=self.initializer(self.target_sparsity))+++  def _validate_block(self, var):+    if self._block_size != [1, 1]:+      if var.get_shape().ndims != 2:+        raise ValueError('Block Sparsity can only be used for layers which '+                          'have 2-dimensional weights.')++  def _random_uniform(self, step, *args, **kwargs):+    """Uniform noise distribution"""+    if self._stateless:+      offset_seed = self._seed_offset + kwargs.get('seed', 0)+      kwargs['seed'] = tf.cast(+        tf.stack([offset_seed, step], tf.int32)+      )+      return tf.random.stateless_uniform(*args, **kwargs)+    return tf.random.uniform(*args, **kwargs)++  def _random_normal(self, step, *args, **kwargs):+    """Gaussian noise distribution"""+    if self._stateless:+      offset_seed = self._seed_offset + kwargs.get('seed', 0)+      kwargs['seed'] = tf.cast(+        tf.stack([offset_seed, step], tf.int32)+      )+      return tf.random.stateless_normal(*args, **kwargs)+    return tf.random.normal(*args, **kwargs)++  def _get_grow_grads(self, optimizer, step, mask, grad):+    """Grow connections baased on gradient information.+    """+    grad_scores = tf.math.abs(grad)+    return grad_scores++  def _get_drop_weights(self, optimizer, step, mask, weight, noise_std=0):+    """+    Drop connections based on weight magnitude.+    """+    masked_weight = mask * weight+    weight_scores = tf.math.abs(masked_weight)+    # add noise to break ties, although the possibility is very low+    if noise_std != 0:+      weight_scores.assign_add(self._random_normal(

++# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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.+# ==============================================================================+"""Helper functions to add support for iterative magnitude/gradient pruning as seen in the RiGL experiments."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++import tensorflow as tf+from absl import logging++from tensorflow.python.ops import summary_ops_v2+from tensorflow.python.summary import summary as summary_ops_v1+from tensorflow_model_optimization.python.core.sparsity.keras import pruning_utils+from tensorflow_model_optimization.python.core.sparsity_tf2 import schedule as update_schedule+from tensorflow_model_optimization.python.core.sparsity_tf2 import pruner+from tensorflow_model_optimization.python.core.sparsity_tf2 import sparse_utils+++class RiGLPruner(pruner.Pruner):+  """+  Implementation of the RiGL dynamic sparse training algorithm.+  """++  def __init__(self,+      update_schedule=update_schedule.ConstantSchedule(0.5, 0),+      sparsity=0.5,+      block_size=(1,1),+      block_pooling_type='AVG',+      initializer=sparse_utils.PermuteOnes,+      stateless=False,+      seed=0,+      seed_offset=0,+      noise_std=0,+      reinit=False+    ):+    """The logic for magnitude-based pruning weight tensors.++    Args:+      update_schedule: A `PruningSchedule` object that controls pruning rate+        throughout training.+      sparsity: the sparsity at which the dynamic sparse training method uses+      block_size: The dimensions (height, weight) for the block sparse pattern+        in rank-2 weight tensors.+      block_pooling_type: (optional) The function to use to pool weights in the+        block. Must be 'AVG' or 'MAX'.+      initializer: the initial sparsity distribution Callable for each layer of the network.+      stateless: whether or not being run on TPU/multi-GPU workers.+      seed: assigned by PruningConfig for multiworker consistency in random processes.+      seed_offset: added to seed to run different experiments.+      reinit: boolean for whether to reinitialize a connection that was drop and regorwn to 

++# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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.+# ==============================================================================+"""Helper functions to add support for iterative magnitude/gradient pruning as seen in the RiGL experiments."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++import tensorflow as tf+from absl import logging++from tensorflow.python.ops import summary_ops_v2+from tensorflow.python.summary import summary as summary_ops_v1+from tensorflow_model_optimization.python.core.sparsity.keras import pruning_utils+from tensorflow_model_optimization.python.core.sparsity_tf2 import schedule as update_schedule+from tensorflow_model_optimization.python.core.sparsity_tf2 import pruner+from tensorflow_model_optimization.python.core.sparsity_tf2 import sparse_utils+++class RiGLPruner(pruner.Pruner):+  """+  Implementation of the RiGL dynamic sparse training algorithm.+  """++  def __init__(self,+      update_schedule=update_schedule.ConstantSchedule(0.5, 0),+      sparsity=0.5,+      block_size=(1,1),+      block_pooling_type='AVG',+      initializer=sparse_utils.PermuteOnes,+      stateless=False,+      seed=0,+      seed_offset=0,+      noise_std=0,+      reinit=False+    ):+    """The logic for magnitude-based pruning weight tensors.++    Args:+      update_schedule: A `PruningSchedule` object that controls pruning rate+        throughout training.+      sparsity: the sparsity at which the dynamic sparse training method uses+      block_size: The dimensions (height, weight) for the block sparse pattern+        in rank-2 weight tensors.+      block_pooling_type: (optional) The function to use to pool weights in the+        block. Must be 'AVG' or 'MAX'.+      initializer: the initial sparsity distribution Callable for each layer of the network.+      stateless: whether or not being run on TPU/multi-GPU workers.+      seed: assigned by PruningConfig for multiworker consistency in random processes.+      seed_offset: added to seed to run different experiments.

++# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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.+# ==============================================================================+"""Helper functions to add support for iterative magnitude/gradient pruning as seen in the RiGL experiments."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++import tensorflow as tf+from absl import logging++from tensorflow.python.ops import summary_ops_v2+from tensorflow.python.summary import summary as summary_ops_v1+from tensorflow_model_optimization.python.core.sparsity.keras import pruning_utils+from tensorflow_model_optimization.python.core.sparsity_tf2 import schedule as update_schedule+from tensorflow_model_optimization.python.core.sparsity_tf2 import pruner+from tensorflow_model_optimization.python.core.sparsity_tf2 import sparse_utils+++class RiGLPruner(pruner.Pruner):+  """+  Implementation of the RiGL dynamic sparse training algorithm.+  """++  def __init__(self,+      update_schedule=update_schedule.ConstantSchedule(0.5, 0),+      sparsity=0.5,+      block_size=(1,1),+      block_pooling_type='AVG',+      initializer=sparse_utils.PermuteOnes,+      stateless=False,+      seed=0,+      seed_offset=0,+      noise_std=0,+      reinit=False+    ):+    """The logic for magnitude-based pruning weight tensors.++    Args:+      update_schedule: A `PruningSchedule` object that controls pruning rate+        throughout training.+      sparsity: the sparsity at which the dynamic sparse training method uses+      block_size: The dimensions (height, weight) for the block sparse pattern+        in rank-2 weight tensors.+      block_pooling_type: (optional) The function to use to pool weights in the+        block. Must be 'AVG' or 'MAX'.+      initializer: the initial sparsity distribution Callable for each layer of the network.+      stateless: whether or not being run on TPU/multi-GPU workers.

++# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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.+# ==============================================================================+"""Helper functions to add support for iterative magnitude/gradient pruning as seen in the RiGL experiments."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++import tensorflow as tf+from absl import logging++from tensorflow.python.ops import summary_ops_v2+from tensorflow.python.summary import summary as summary_ops_v1+from tensorflow_model_optimization.python.core.sparsity.keras import pruning_utils+from tensorflow_model_optimization.python.core.sparsity_tf2 import schedule as update_schedule+from tensorflow_model_optimization.python.core.sparsity_tf2 import pruner+from tensorflow_model_optimization.python.core.sparsity_tf2 import sparse_utils+++class RiGLPruner(pruner.Pruner):+  """+  Implementation of the RiGL dynamic sparse training algorithm.+  """++  def __init__(self,+      update_schedule=update_schedule.ConstantSchedule(0.5, 0),+      sparsity=0.5,+      block_size=(1,1),+      block_pooling_type='AVG',+      initializer=sparse_utils.PermuteOnes,

+# Extension Types++| Status        | Proposed                                             |+:-------------- |:---------------------------------------------------- |+| **Authors**   | Edward Loper ([email protected]) |+| **Sponsor**   | Alex Passos ([email protected])                     |+| **Updated**   | 2020-07-21                                           |++## Objective++This RFC proposes a protocol that can be used to define **_user-defined+object-oriented Python types_** that are supported by TensorFlow APIs.++## Motivation++Object oriented types can make systems more readable, modular, maintainable.+However, most TensorFlow APIs do not currently support user-defined Python+types.  This includes both high-level APIs (such as `Keras`, `tf.function`,+`tf.SavedModel`) and lower-level APIs (such as `tf.while_loop` and `tf.add`).  ++This RFC proposes a set of protocols that will allow TensorFlow APIs to handle+user-defined Python types.  A version of this interface is already used+internally to implement several core TensorFlow data types, including+`tf.SparseTensor`, `tf.RaggedTensor`, `tf.data.Dataset`, and+`tf.StructuredTensor`.++At a high level, types supported by this interface can be divided into two broad+categories:++* **_General data structures_**.+  These types are handled by "generic" APIs whose behavior does not depend on+  the value of each object (such as `tf.function`, `SavedModel`, and+  `tf.while_loop`).++* **_Tensor-like types_**, which specialize or extend tf.Tensor.+  Values of these types have a `rank`, a `shape`, and usually a `dtype`.  In+  addition to the "generic" APIs, these types can be handled by Tensor-specific+  APIs (such as `tf.stack`, `tf.add`, and `tf.reduce_mean`).++Examples of user-defined types that could defined or extended with this protocol+include:++**General data structures:**++* `tfp.Distribution`: Encodes a statistical distribution.+* `TensorDigraph`: Encodes the set of nodes and edges in a directed graph.+* `DimensionAlignment`: Encodes a correspondence between two related dimensions+  (e.g., between a `word` dimension and a `speaker` dimension).++**Tensor-like types:**++* `CSRSparseTensor`: A sparsely-encoded tensor that uses the Compressed Sparse+  Row encoding.+* `MaskedTensor`: Pairs a Tensor with a corresponding boolean mask, indicating+  which values are valid, and automatically updates the mask as appropriate+  when used with TensorFlow ops (such as `tf.add` or `tf.reduce_sum`).+* `LabeledTensor`: Pairs a Tensor with a list of axis names, which can be used+  for error detection and reporting, for specifying axes by name, and for+  broadcasting.++## User Benefit++### Object-Oriented TensorFlow++This proposal brings the benefits of Object-Oriented Programming to TensorFlow+users, allowing them to define modular encapsulated data structures that+interoperate with TensorFlow APIs.  This allows TensorFlow models to be defined+at a higher level of abstraction.++### Development outside the main TensorFlow repository++Prior to this proposal, the only way to develop such data structures (e.g.+`SparseTensor`) was to develop them inside the main TensorFlow code base.  This+introduced significant barriers to rapid development, including slow release+cycles, strong backwards compatibility constraints, and centralized API+approval.  By allowing such data structures to be developed outside the main+TensorFlow code base, we hope to make it much easier to experiment with new+types and designs.  If general-purpose types are developed that become+sufficiently mature, we may consider bringing them into the main TensorFlow code+base.++### TensorFlow APIs supported by user-defined types++User-defined types that implement the interface proposed by this RFC will be+supported by the following APIs:++* **Keras**: User-defined types can be used as inputs and outputs for Keras+  `Models` and `Layers`.

 def _create_slots(self, var_list):         pruner.create_slots(self, var)    def _resource_apply_dense(self, grad, var, apply_state):+    self.preprocess_weights(self, var, grad)

+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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 key functions in lthpruner library."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++# import g3++from absl.testing import parameterized+import numpy as np+import tensorflow as tf+print("Running TF", tf.__version__)++# TODO(b/139939526): move to public API.+from tensorflow.python.keras import keras_parameterized+from tensorflow_model_optimization.python.core.keras import compat+from tensorflow_model_optimization.python.core.sparsity.keras import pruning_schedule+from tensorflow_model_optimization.python.core.sparsity.keras import pruning_utils+from tensorflow_model_optimization.python.core.sparsity_tf2 import lthpruner as pruner++dtypes = tf.dtypes+test = tf.test++def get_lth_sparsity(save_round, n_rounds, target_sparsity, end_epoch):+  """+  save_round: when to save the weight; 1 onwards+  n_rounds: number of pruning cycles to do+  target_sparsity: the sparsity percentage to achieve by the end of the iteration+  end_epoch: total number of epochs to train for where pruning is eligible++  Returns:+    percent to prune to after each cyle/round+  """+  # no pruning until weights are saved TODO: off by one??+  n_rounds = tf.constant(n_rounds, dtype='float32') # dtype='float32'+  frequency = tf.math.floordiv(end_epoch - save_round + 1, n_rounds) # range(0, end, freq)+  prune_ratio_per_round = tf.math.pow(target_sparsity, tf.math.divide(1, n_rounds))+  return tf.cast(frequency, tf.int64), prune_ratio_per_round++def make_pruning_schedule(target_sparsity, begin, end, freq):+  return pruning_schedule.ConstantSparsity(target_sparsity, begin, end, freq)++def sample_noise(x, mu=0, sigma=1.):+  sample = tf.random.normal((), mean=mu,  stddev=sigma, dtype=tf.float64)+  return sample++def _dummy_gradient(x, dtype=tf.float32):+  try:+    base_type = x.dtype+  except:+    base_type = dtype+  grad = tf.ones_like(x, dtype=base_type)+  return grad++class PruningTest(test.TestCase, parameterized.TestCase):++  def setUp(self):+    super(PruningTest, self).setUp()+    self.block_size = (1, 1)+    self.block_pooling_type = "AVG"+    self.target_sparsity = 0.5+    self.constant_sparsity = pruning_schedule.ConstantSparsity(self.target_sparsity, 0, 100, 1)+    self.save_init = tf.Variable(0)+    self.save_itr_10 = 10+    self.dummy_optimizer = tf.keras.optimizers.SGD(learning_rate=0.01)+    self.grad = _dummy_gradient++  # Variable initialization outside of setUp() is needed for compatibility with+  # run_all_keras_modes.+  #+  # setUp() lies outside of the "eager scope" that wraps the test cases+  # themselves, resulting in initializing graph tensors instead of eager+  # tensors when testing eager execution.++  def testUpdateSingleMask(self):+    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")+    weight_dtype = weight.dtype.base_dtype+    mask = tf.Variable(+        tf.ones(weight.get_shape(), dtype=weight_dtype),+        name="mask",+        dtype=weight_dtype)+    threshold = tf.Variable(+        tf.zeros([], dtype=weight_dtype), name="threshold", dtype=weight_dtype)+    pruning_vars = [(weight, mask, threshold)]++    p = pruner.LTHPruner(+        pruning_schedule=self.constant_sparsity,+        save_iteration=self.save_init,+        block_size=self.block_size,+        block_pooling_type=self.block_pooling_type)++    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)+    p.create_slots(optimizer, weight)++    mask_before_pruning = optimizer.get_slot(weight, "mask").read_value()+    self.assertAllEqual(np.count_nonzero(mask_before_pruning), 100)++    next_step = optimizer.iterations.assign_add(1)+    p.update_masks(pruning_vars, next_step)++    mask_after_pruning = mask.read_value()+    self.assertAllEqual(np.count_nonzero(mask_after_pruning), 50)++  def testConstructsMaskAndThresholdCorrectly(self):+    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")+    weight_dtype = weight.dtype.base_dtype+    p = pruner.LTHPruner(+        # Sparsity math often returns values with small tolerances.+        pruning_schedule=lambda x: (True, 0.200000018),+        save_iteration=self.save_init,+        block_size=(1, 1), block_pooling_type=None)+        +    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)+    p.create_slots(optimizer, weight)+    step = optimizer.iterations++    # input matrix is [ 1.0, 2.0, ..., 8.0, 9.0, 10.0 ]+    threshold, mask = p._update_mask(step, np.arange(1, 11))++    self.assertEqual(3, threshold)+    self.assertAllEqual(+        # expected matrix is [ 0.0, 0.0, 1.0, 1.0 ... 1.0 ]+        np.concatenate((np.zeros(2), np.ones(8))), mask)++  def _blockMasking(self, block_size, block_pooling_type, weight,+                    expected_mask):+    mask = tf.Variable(+        tf.ones(weight.get_shape(), dtype=weight.dtype),+        name="mask",+        dtype=weight.dtype)+    threshold = tf.Variable(+        tf.zeros([], dtype=weight.dtype), name="threshold", dtype=weight.dtype)++    # Set up pruning+    p = pruner.LTHPruner(+        pruning_schedule=self.constant_sparsity,+        save_iteration=self.save_init,+        block_size=block_size,+        block_pooling_type=block_pooling_type)++    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)+    p.create_slots(optimizer, weight)+    step = optimizer.iterations++    _, new_mask = p._maybe_update_block_mask(step, weight)+    # Check if the mask is the same size as the weights+    self.assertAllEqual(new_mask.get_shape(), weight.get_shape())+    mask_after_pruning = new_mask+    self.assertAllEqual(mask_after_pruning, expected_mask)++  def testBlockMaskingAvg(self):+    block_size = (2, 2)+    block_pooling_type = "AVG"+    weight = tf.Variable([[0.1, 0.1, 0.2, 0.2], [0.1, 0.1, 0.2, 0.2],+                          [0.3, 0.3, 0.4, 0.4], [0.3, 0.3, 0.4, 0.4]])+    expected_mask = [[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0],+                     [1., 1., 1., 1.], [1., 1., 1., 1.]]++    self._blockMasking(block_size, block_pooling_type, weight, expected_mask)++  def testBlockMaskingMax(self):+    block_size = (2, 2)+    block_pooling_type = "MAX"+    weight = tf.Variable([[0.1, 0.0, 0.2, 0.0], [0.0, -0.1, 0.0, -0.2],+                                   [0.3, 0.0, 0.4, 0.0], [0.0, -0.3, 0.0,+                                                          -0.4]])+    expected_mask = [[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0],+                     [1., 1., 1., 1.], [1., 1., 1., 1.]]++    self._blockMasking(block_size, block_pooling_type, weight, expected_mask)++  def testBlockMaskingWithHigherDimensionsRaisesError(self):+    block_size = (2, 2)+    block_pooling_type = "AVG"+    # Weights as in testBlockMasking, but with one extra dimension.+    weight = tf.Variable([[[0.1, 0.1, 0.2, 0.2], [0.1, 0.1, 0.2, 0.2],+                                    [0.3, 0.3, 0.4, 0.4], [0.3, 0.3, 0.4,+                                                           0.4]]])+    expected_mask = [[[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0],+                      [1., 1., 1., 1.], [1., 1., 1., 1.]]]++    # Block masking should only be used with 2 Dimensional weights.+    with self.assertRaises(ValueError):+      self._blockMasking(block_size, block_pooling_type, weight, expected_mask)++  def testConditionalMaskUpdate(self):+    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")+    weight_dtype = weight.dtype.base_dtype+    mask = tf.Variable(+        tf.ones(weight.get_shape(), dtype=weight_dtype),+        name="mask",+        dtype=weight_dtype)+    threshold = tf.Variable(+        tf.zeros([], dtype=weight_dtype), name="threshold", dtype=weight_dtype)+    pruning_vars = [(weight, mask, threshold)]++    def linear_sparsity(step):+      sparsity_val = tf.convert_to_tensor(+          [0.0, 0.1, 0.1, 0.3, 0.3, 0.5, 0.5, 0.5, 0.5, 0.5])+      return tf.convert_to_tensor(True), sparsity_val[step]++    def weight_mask_op(pruning_vars):+      values_and_vars = []+      for weight, mask, _ in  pruning_vars:+        weight.assign(tf.math.multiply(weight, mask))++    # Set up pruning+    p = pruner.LTHPruner(+        pruning_schedule=linear_sparsity,+        save_iteration=self.save_init,+        block_size=self.block_size,+        block_pooling_type=self.block_pooling_type)++    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)+    p.create_slots(optimizer, weight)+    step = optimizer.iterations++    non_zero_count = []+    for _ in range(10):+      step = optimizer.iterations+      p.update_masks(pruning_vars, step)+      weight_mask_op(pruning_vars)+      optimizer.iterations.assign_add(1)++      non_zero_count.append(tf.math.count_nonzero(weight))++    # Weights pruned at steps 1,3,5+    expected_non_zero_count = [100, 90, 90, 70, 70, 50, 50, 50, 50, 50]+    self.assertAllEqual(expected_non_zero_count, non_zero_count)+++  def testSaveOriginalInitializations(self):+    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")+    weight_dtype = weight.dtype.base_dtype+    save_round = 0+    n_rounds = 5+    end_iter = 100+    frequency, prune_ratio_per_round = get_lth_sparsity(save_round, n_rounds, self.target_sparsity, end_iter)+    pruning_schedule = make_pruning_schedule(prune_ratio_per_round, 0, end_iter, frequency)+    +    p = pruner.LTHPruner(+        pruning_schedule=pruning_schedule,+        save_iteration=self.save_init,+        block_size=self.block_size,+        block_pooling_type=self.block_pooling_type)++    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)+    init_weights_before_pruning = weight.read_value()++    def _train():+      p.create_slots(optimizer, weight)+      p.preprocess_weights(optimizer, weight, self.grad(weight))+      weight.assign(tf.math.add(weight, sample_noise(0))) # update weights+      p.postprocess_weights(optimizer, weight, self.grad(weight))+      optimizer.iterations.assign_add(1)+      initialization_slot = optimizer.get_slot(weight, "original_initialization")+      return initialization_slot+    +    initialization_slot = _train()+    self.assertAllEqual(initialization_slot, init_weights_before_pruning)+    +    initialization_slot = tf.function(_train)()+    self.assertAllEqual(initialization_slot, init_weights_before_pruning)+++  def testSaveWeightsIterK(self):+    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")++    save_round = 5+    n_rounds = 24+    end_iter = 100+    frequency, prune_ratio_per_round = get_lth_sparsity(save_round, n_rounds, self.target_sparsity, end_iter)+    pruning_schedule = make_pruning_schedule(1 - prune_ratio_per_round, save_round, end_iter, frequency)+    +    p = pruner.LTHPruner(+        pruning_schedule=pruning_schedule,+        save_iteration=save_round,+        block_size=self.block_size,+        block_pooling_type=self.block_pooling_type)++    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)++    def _train(weight):+      expected_saved_initialization = tf.ones_like(weight, dtype=weight.dtype.base_dtype) * -1+      p.create_slots(optimizer, weight)+      for i in range(7):+        p.preprocess_weights(optimizer, weight, self.grad(weight))+        if optimizer.iterations == save_round:+          expected_saved_initialization = weight.read_value()+        weight.assign(tf.math.add(weight, sample_noise(i)))+        p.postprocess_weights(optimizer, weight, self.grad(weight))+        optimizer.iterations.assign_add(1)+      return expected_saved_initialization+      +    expected_saved_initialization = _train(weight)+    self.assertAllEqual(optimizer.get_slot(weight, "original_initialization"), expected_saved_initialization)++    initialization_slot_k = tf.math.multiply(optimizer.get_slot(weight, "original_initialization"), optimizer.get_slot(weight, "mask"))+    masked_expected = tf.math.multiply(expected_saved_initialization, optimizer.get_slot(weight, "mask"))+    self.assertAllEqual(initialization_slot_k, masked_expected)++    mask_after_pruning = optimizer.get_slot(weight, "mask").read_value()+    self.assertAllEqual(np.count_nonzero(mask_after_pruning), 97)++    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")+    optimizer.iterations.assign(0)+    expected_saved_initialization = tf.function(_train)(weight)+    self.assertAllEqual(optimizer.get_slot(weight, "original_initialization"), expected_saved_initialization)++    initialization_slot_k = tf.math.multiply(optimizer.get_slot(weight, "original_initialization"), optimizer.get_slot(weight, "mask"))+    masked_expected = tf.math.multiply(expected_saved_initialization, optimizer.get_slot(weight, "mask"))+    self.assertAllEqual(initialization_slot_k, masked_expected)++    mask_after_pruning = optimizer.get_slot(weight, "mask").read_value()+    self.assertAllEqual(np.count_nonzero(mask_after_pruning), 97)++  def testReloadWeightsatInitialization(self):+    weight1 = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")++    save_round = 0+    n_rounds = 20+    end_iter = 100+    frequency, prune_ratio_per_round = get_lth_sparsity(save_round, n_rounds, self.target_sparsity, end_iter)+    pruning_schedule = make_pruning_schedule(1 - prune_ratio_per_round, save_round, end_iter, frequency)+    +    p = pruner.LTHPruner(+        pruning_schedule=pruning_schedule,+        save_iteration=self.save_init,+        block_size=self.block_size,+        block_pooling_type=self.block_pooling_type)++    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)+    expected_saved_initialization = None++    def _train(weight):+      expected_saved_initialization = tf.ones_like(weight, dtype=weight.dtype.base_dtype) * -1+      p.create_slots(optimizer, weight)+      for i in tf.range(0, 2):+        p.preprocess_weights(optimizer, weight, self.grad(weight))+        if i == save_round:+          expected_saved_initialization = weight.read_value()+        weight.assign(tf.math.add(weight, sample_noise(i)))+        p.postprocess_weights(optimizer, weight, self.grad(weight))+        optimizer.iterations.assign_add(1) # save weights right before iteration+      return expected_saved_initialization++    expected_saved_initialization = _train(weight1)+    initialization_slot = optimizer.get_slot(weight1, "original_initialization")+    self.assertAllEqual(initialization_slot, expected_saved_initialization)++    mask_after_pruning = optimizer.get_slot(weight1, "mask").read_value()+    masked_weight_expected = tf.math.multiply(initialization_slot, mask_after_pruning)+    self.assertAllEqual(np.count_nonzero(masked_weight_expected), 97)+    self.assertAllEqual(np.count_nonzero(mask_after_pruning), 97)++    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")+    optimizer.iterations.assign(0)+    expected_saved_initialization = tf.function(_train)(weight)+    initialization_slot = optimizer.get_slot(weight, "original_initialization")+    self.assertAllEqual(initialization_slot, expected_saved_initialization)++    mask_after_pruning = optimizer.get_slot(weight, "mask").read_value()+    masked_weight_expected = tf.math.multiply(initialization_slot, mask_after_pruning)+    self.assertAllEqual(np.count_nonzero(masked_weight_expected), 97)+    self.assertAllEqual(np.count_nonzero(mask_after_pruning), 97)++  +  def testReloadWeightsIterationK(self):+    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")++    save_round = 5+    n_rounds = 24+    end_iter = 100+    frequency, prune_ratio_per_round = get_lth_sparsity(save_round, n_rounds, self.target_sparsity, end_iter)+    pruning_schedule = make_pruning_schedule(1 - prune_ratio_per_round, save_round + 1, end_iter, frequency)+    +    p = pruner.LTHPruner(+        pruning_schedule=pruning_schedule,+        save_iteration=save_round,+        block_size=self.block_size,+        block_pooling_type=self.block_pooling_type)++    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)++    def _train(weight):+      p.create_slots(optimizer, weight)+      for i in range(7):+        p.preprocess_weights(optimizer, weight, self.grad(weight))+        weight.assign(tf.math.add(weight, sample_noise(i)))+        p.postprocess_weights(optimizer, weight, self.grad(weight))+        optimizer.iterations.assign_add(1)++    _train(weight)+    masked_orig_init = tf.math.multiply(optimizer.get_slot(weight, "original_initialization"), optimizer.get_slot(weight, 'mask'))+    self.assertAllEqual(masked_orig_init, weight)++    mask_after_pruning = optimizer.get_slot(weight, "mask").read_value()+    self.assertAllEqual(np.count_nonzero(mask_after_pruning), 97)++    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")+    optimizer.iterations.assign(0)+    tf.function(_train)(weight)+    masked_orig_init = tf.math.multiply(optimizer.get_slot(weight, "original_initialization"), optimizer.get_slot(weight, 'mask'))+    self.assertAllEqual(masked_orig_init, weight)++    mask_after_pruning = optimizer.get_slot(weight, "mask").read_value()+    self.assertAllEqual(np.count_nonzero(mask_after_pruning), 97)+++  def testReloadTwoTimes(self):+    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")+    weight_dtype = weight.dtype.base_dtype++    save_round = 5+    n_rounds = 24+    end_iter = 100+    frequency, prune_ratio_per_round = get_lth_sparsity(save_round, n_rounds, self.target_sparsity, end_iter)+    pruning_schedule = make_pruning_schedule(1 - prune_ratio_per_round, save_round + 1, end_iter, frequency)+    +    p = pruner.LTHPruner(+        pruning_schedule=pruning_schedule,+        save_iteration=save_round,+        block_size=self.block_size,+        block_pooling_type=self.block_pooling_type)++    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)++    def _train(weight):+      expected_first_saved_initialization = None+      expected_second_saved_initialization = None+      p.create_slots(optimizer, weight)+      for i in range(0, 13): # this should save, reload, and update exactly twice+        p.preprocess_weights(optimizer, weight, self.grad(weight))+        if i == save_round - 1:+          expected_first_saved_initialization = optimizer.get_slot(weight, "original_initialization")+        if i == (save_round - 1) * 2:+          expected_second_saved_initialization = optimizer.get_slot(weight, "original_initialization")+        weight.assign(tf.math.add(weight, sample_noise(i)))+        p.postprocess_weights(optimizer, weight, self.grad(weight))+        optimizer.iterations.assign_add(1)+      return expected_first_saved_initialization, expected_second_saved_initialization++    expected_first_saved_initialization, expected_second_saved_initialization = _train(weight)++    self.assertAllEqual(expected_first_saved_initialization, expected_second_saved_initialization)++    initialization_slot = optimizer.get_slot(weight, "original_initialization")+    self.assertAllEqual(initialization_slot, expected_first_saved_initialization)

 def _create_slots(self, var_list):         pruner.create_slots(self, var)    def _resource_apply_dense(self, grad, var, apply_state):+    self.preprocess_weights(self, var, grad)     self._optimizer._resource_apply_dense(grad, var, apply_state)-    self.prune(var, grad)+    self.postprocess_weights(self, var, grad)    def _resource_apply_sparse(self, grad, var, indices, **kwargs):+    self.preprocess_weights(self, var, grad)

 def _create_slots(self, var_list):         pruner.create_slots(self, var)    def _resource_apply_dense(self, grad, var, apply_state):+    self.preprocess_weights(self, var, grad)

 def testReloadWeightsatInitialization(self):         block_size=self.block_size,         block_pooling_type=self.block_pooling_type) -    optimizer = self.dummy_optimizer+    optimizer =   tf.keras.optimizers.SGD(learning_rate=0.01)     optimizer.iterations.assign(0)     expected_saved_initialization = None -    def _train(weight):+    def _train(optimizer, weight):       expected_saved_initialization = tf.ones_like(weight, dtype=weight.dtype.base_dtype) * -1+      expected_reload = tf.ones_like(weight, dtype=weight.dtype.base_dtype) * -1       p.create_slots(optimizer, weight)       for i in tf.range(0, 2):         p.preprocess_weights(optimizer, weight, self.grad(weight))         if i == save_round:           expected_saved_initialization = weight.read_value()         weight.assign(tf.math.add(weight, sample_noise(i)))         p.postprocess_weights(optimizer, weight, self.grad(weight))+        should_reload = p._reload_schedule._should_prune_in_step(optimizer.iterations,

+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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 key functions in lthpruner library."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++# import g3++from absl.testing import parameterized+import numpy as np+import tensorflow as tf+print("Running TF", tf.__version__)++# TODO(b/139939526): move to public API.+from tensorflow.python.keras import keras_parameterized+from tensorflow_model_optimization.python.core.keras import compat+from tensorflow_model_optimization.python.core.sparsity.keras import pruning_schedule+from tensorflow_model_optimization.python.core.sparsity.keras import pruning_utils+from tensorflow_model_optimization.python.core.sparsity_tf2 import lthpruner as pruner++dtypes = tf.dtypes+test = tf.test++def get_lth_sparsity(save_round, n_rounds, target_sparsity, end_epoch):+  """+  save_round: when to save the weight; 1 onwards+  n_rounds: number of pruning cycles to do+  target_sparsity: the sparsity percentage to achieve by the end of the iteration+  end_epoch: total number of epochs to train for where pruning is eligible++  Returns:+    percent to prune to after each cyle/round+  """+  # no pruning until weights are saved TODO: off by one??+  n_rounds = tf.constant(n_rounds, dtype='float32') # dtype='float32'+  frequency = tf.math.floordiv(end_epoch - save_round + 1, n_rounds) # range(0, end, freq)+  prune_ratio_per_round = tf.math.pow(target_sparsity, tf.math.divide(1, n_rounds))+  return tf.cast(frequency, tf.int64), prune_ratio_per_round++def make_pruning_schedule(target_sparsity, begin, end, freq):+  return pruning_schedule.ConstantSparsity(target_sparsity, begin, end, freq)++def sample_noise(x, mu=0, sigma=1.):+  sample = tf.random.normal((), mean=mu,  stddev=sigma, dtype=tf.float64)+  return sample++def _dummy_gradient(x, dtype=tf.float32):+  try:+    base_type = x.dtype+  except:+    base_type = dtype+  grad = tf.ones_like(x, dtype=base_type)+  return grad++class PruningTest(test.TestCase, parameterized.TestCase):++  def setUp(self):+    super(PruningTest, self).setUp()+    self.block_size = (1, 1)+    self.block_pooling_type = "AVG"+    self.target_sparsity = 0.5+    self.constant_sparsity = pruning_schedule.ConstantSparsity(self.target_sparsity, 0, 100, 1)+    self.save_init = tf.Variable(0)+    self.save_itr_10 = 10+    self.dummy_optimizer = tf.keras.optimizers.SGD(learning_rate=0.01)+    self.grad = _dummy_gradient++  # Variable initialization outside of setUp() is needed for compatibility with+  # run_all_keras_modes.+  #+  # setUp() lies outside of the "eager scope" that wraps the test cases+  # themselves, resulting in initializing graph tensors instead of eager+  # tensors when testing eager execution.++  def testUpdateSingleMask(self):+    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")+    weight_dtype = weight.dtype.base_dtype+    mask = tf.Variable(+        tf.ones(weight.get_shape(), dtype=weight_dtype),+        name="mask",+        dtype=weight_dtype)+    threshold = tf.Variable(+        tf.zeros([], dtype=weight_dtype), name="threshold", dtype=weight_dtype)+    pruning_vars = [(weight, mask, threshold)]++    p = pruner.LTHPruner(+        pruning_schedule=self.constant_sparsity,+        save_iteration=self.save_init,+        block_size=self.block_size,+        block_pooling_type=self.block_pooling_type)++    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)+    p.create_slots(optimizer, weight)++    mask_before_pruning = optimizer.get_slot(weight, "mask").read_value()+    self.assertAllEqual(np.count_nonzero(mask_before_pruning), 100)++    next_step = optimizer.iterations.assign_add(1)+    p.update_masks(pruning_vars, next_step)++    mask_after_pruning = mask.read_value()+    self.assertAllEqual(np.count_nonzero(mask_after_pruning), 50)++  def testConstructsMaskAndThresholdCorrectly(self):+    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")+    weight_dtype = weight.dtype.base_dtype+    p = pruner.LTHPruner(+        # Sparsity math often returns values with small tolerances.+        pruning_schedule=lambda x: (True, 0.200000018),+        save_iteration=self.save_init,+        block_size=(1, 1), block_pooling_type=None)+        +    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)+    p.create_slots(optimizer, weight)+    step = optimizer.iterations++    # input matrix is [ 1.0, 2.0, ..., 8.0, 9.0, 10.0 ]+    threshold, mask = p._update_mask(step, np.arange(1, 11))++    self.assertEqual(3, threshold)+    self.assertAllEqual(+        # expected matrix is [ 0.0, 0.0, 1.0, 1.0 ... 1.0 ]+        np.concatenate((np.zeros(2), np.ones(8))), mask)++  def _blockMasking(self, block_size, block_pooling_type, weight,+                    expected_mask):+    mask = tf.Variable(+        tf.ones(weight.get_shape(), dtype=weight.dtype),+        name="mask",+        dtype=weight.dtype)+    threshold = tf.Variable(+        tf.zeros([], dtype=weight.dtype), name="threshold", dtype=weight.dtype)++    # Set up pruning+    p = pruner.LTHPruner(+        pruning_schedule=self.constant_sparsity,+        save_iteration=self.save_init,+        block_size=block_size,+        block_pooling_type=block_pooling_type)++    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)+    p.create_slots(optimizer, weight)+    step = optimizer.iterations++    _, new_mask = p._maybe_update_block_mask(step, weight)+    # Check if the mask is the same size as the weights+    self.assertAllEqual(new_mask.get_shape(), weight.get_shape())+    mask_after_pruning = new_mask+    self.assertAllEqual(mask_after_pruning, expected_mask)++  def testBlockMaskingAvg(self):+    block_size = (2, 2)+    block_pooling_type = "AVG"+    weight = tf.Variable([[0.1, 0.1, 0.2, 0.2], [0.1, 0.1, 0.2, 0.2],+                          [0.3, 0.3, 0.4, 0.4], [0.3, 0.3, 0.4, 0.4]])+    expected_mask = [[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0],+                     [1., 1., 1., 1.], [1., 1., 1., 1.]]++    self._blockMasking(block_size, block_pooling_type, weight, expected_mask)++  def testBlockMaskingMax(self):+    block_size = (2, 2)+    block_pooling_type = "MAX"+    weight = tf.Variable([[0.1, 0.0, 0.2, 0.0], [0.0, -0.1, 0.0, -0.2],+                                   [0.3, 0.0, 0.4, 0.0], [0.0, -0.3, 0.0,+                                                          -0.4]])+    expected_mask = [[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0],+                     [1., 1., 1., 1.], [1., 1., 1., 1.]]++    self._blockMasking(block_size, block_pooling_type, weight, expected_mask)++  def testBlockMaskingWithHigherDimensionsRaisesError(self):+    block_size = (2, 2)+    block_pooling_type = "AVG"+    # Weights as in testBlockMasking, but with one extra dimension.+    weight = tf.Variable([[[0.1, 0.1, 0.2, 0.2], [0.1, 0.1, 0.2, 0.2],+                                    [0.3, 0.3, 0.4, 0.4], [0.3, 0.3, 0.4,+                                                           0.4]]])+    expected_mask = [[[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0],+                      [1., 1., 1., 1.], [1., 1., 1., 1.]]]++    # Block masking should only be used with 2 Dimensional weights.+    with self.assertRaises(ValueError):+      self._blockMasking(block_size, block_pooling_type, weight, expected_mask)++  def testConditionalMaskUpdate(self):+    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")+    weight_dtype = weight.dtype.base_dtype+    mask = tf.Variable(+        tf.ones(weight.get_shape(), dtype=weight_dtype),+        name="mask",+        dtype=weight_dtype)+    threshold = tf.Variable(+        tf.zeros([], dtype=weight_dtype), name="threshold", dtype=weight_dtype)+    pruning_vars = [(weight, mask, threshold)]++    def linear_sparsity(step):+      sparsity_val = tf.convert_to_tensor(+          [0.0, 0.1, 0.1, 0.3, 0.3, 0.5, 0.5, 0.5, 0.5, 0.5])+      return tf.convert_to_tensor(True), sparsity_val[step]++    def weight_mask_op(pruning_vars):+      values_and_vars = []+      for weight, mask, _ in  pruning_vars:+        weight.assign(tf.math.multiply(weight, mask))++    # Set up pruning+    p = pruner.LTHPruner(+        pruning_schedule=linear_sparsity,+        save_iteration=self.save_init,+        block_size=self.block_size,+        block_pooling_type=self.block_pooling_type)++    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)+    p.create_slots(optimizer, weight)+    step = optimizer.iterations++    non_zero_count = []+    for _ in range(10):+      step = optimizer.iterations+      p.update_masks(pruning_vars, step)+      weight_mask_op(pruning_vars)+      optimizer.iterations.assign_add(1)++      non_zero_count.append(tf.math.count_nonzero(weight))++    # Weights pruned at steps 1,3,5+    expected_non_zero_count = [100, 90, 90, 70, 70, 50, 50, 50, 50, 50]+    self.assertAllEqual(expected_non_zero_count, non_zero_count)+++  def testSaveOriginalInitializations(self):+    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")+    weight_dtype = weight.dtype.base_dtype+    save_round = 0+    n_rounds = 5+    end_iter = 100+    frequency, prune_ratio_per_round = get_lth_sparsity(save_round, n_rounds, self.target_sparsity, end_iter)+    pruning_schedule = make_pruning_schedule(prune_ratio_per_round, 0, end_iter, frequency)+    +    p = pruner.LTHPruner(+        pruning_schedule=pruning_schedule,+        save_iteration=self.save_init,+        block_size=self.block_size,+        block_pooling_type=self.block_pooling_type)++    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)+    init_weights_before_pruning = weight.read_value()++    def _train():+      p.create_slots(optimizer, weight)+      p.preprocess_weights(optimizer, weight, self.grad(weight))+      weight.assign(tf.math.add(weight, sample_noise(0))) # update weights+      p.postprocess_weights(optimizer, weight, self.grad(weight))+      optimizer.iterations.assign_add(1)+      initialization_slot = optimizer.get_slot(weight, "original_initialization")+      return initialization_slot+    +    initialization_slot = _train()+    self.assertAllEqual(initialization_slot, init_weights_before_pruning)+    +    initialization_slot = tf.function(_train)()+    self.assertAllEqual(initialization_slot, init_weights_before_pruning)+++  def testSaveWeightsIterK(self):+    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")++    save_round = 5+    n_rounds = 24+    end_iter = 100+    frequency, prune_ratio_per_round = get_lth_sparsity(save_round, n_rounds, self.target_sparsity, end_iter)+    pruning_schedule = make_pruning_schedule(1 - prune_ratio_per_round, save_round, end_iter, frequency)+    +    p = pruner.LTHPruner(+        pruning_schedule=pruning_schedule,+        save_iteration=save_round,+        block_size=self.block_size,+        block_pooling_type=self.block_pooling_type)++    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)++    def _train(weight):+      expected_saved_initialization = tf.ones_like(weight, dtype=weight.dtype.base_dtype) * -1+      p.create_slots(optimizer, weight)+      for i in range(7):+        p.preprocess_weights(optimizer, weight, self.grad(weight))+        if optimizer.iterations == save_round:+          expected_saved_initialization = weight.read_value()+        weight.assign(tf.math.add(weight, sample_noise(i)))+        p.postprocess_weights(optimizer, weight, self.grad(weight))+        optimizer.iterations.assign_add(1)+      return expected_saved_initialization+      +    expected_saved_initialization = _train(weight)+    self.assertAllEqual(optimizer.get_slot(weight, "original_initialization"), expected_saved_initialization)++    initialization_slot_k = tf.math.multiply(optimizer.get_slot(weight, "original_initialization"), optimizer.get_slot(weight, "mask"))+    masked_expected = tf.math.multiply(expected_saved_initialization, optimizer.get_slot(weight, "mask"))+    self.assertAllEqual(initialization_slot_k, masked_expected)++    mask_after_pruning = optimizer.get_slot(weight, "mask").read_value()+    self.assertAllEqual(np.count_nonzero(mask_after_pruning), 97)++    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")+    optimizer.iterations.assign(0)+    expected_saved_initialization = tf.function(_train)(weight)+    self.assertAllEqual(optimizer.get_slot(weight, "original_initialization"), expected_saved_initialization)++    initialization_slot_k = tf.math.multiply(optimizer.get_slot(weight, "original_initialization"), optimizer.get_slot(weight, "mask"))+    masked_expected = tf.math.multiply(expected_saved_initialization, optimizer.get_slot(weight, "mask"))+    self.assertAllEqual(initialization_slot_k, masked_expected)++    mask_after_pruning = optimizer.get_slot(weight, "mask").read_value()+    self.assertAllEqual(np.count_nonzero(mask_after_pruning), 97)++  def testReloadWeightsatInitialization(self):+    weight1 = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")++    save_round = 0+    n_rounds = 20+    end_iter = 100+    frequency, prune_ratio_per_round = get_lth_sparsity(save_round, n_rounds, self.target_sparsity, end_iter)+    pruning_schedule = make_pruning_schedule(1 - prune_ratio_per_round, save_round, end_iter, frequency)+    +    p = pruner.LTHPruner(+        pruning_schedule=pruning_schedule,+        save_iteration=self.save_init,+        block_size=self.block_size,+        block_pooling_type=self.block_pooling_type)++    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)+    expected_saved_initialization = None++    def _train(weight):+      expected_saved_initialization = tf.ones_like(weight, dtype=weight.dtype.base_dtype) * -1+      p.create_slots(optimizer, weight)+      for i in tf.range(0, 2):+        p.preprocess_weights(optimizer, weight, self.grad(weight))+        if i == save_round:+          expected_saved_initialization = weight.read_value()+        weight.assign(tf.math.add(weight, sample_noise(i)))+        p.postprocess_weights(optimizer, weight, self.grad(weight))+        optimizer.iterations.assign_add(1) # save weights right before iteration+      return expected_saved_initialization++    expected_saved_initialization = _train(weight1)+    initialization_slot = optimizer.get_slot(weight1, "original_initialization")+    self.assertAllEqual(initialization_slot, expected_saved_initialization)++    mask_after_pruning = optimizer.get_slot(weight1, "mask").read_value()+    masked_weight_expected = tf.math.multiply(initialization_slot, mask_after_pruning)+    self.assertAllEqual(np.count_nonzero(masked_weight_expected), 97)+    self.assertAllEqual(np.count_nonzero(mask_after_pruning), 97)++    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")+    optimizer.iterations.assign(0)

+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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 key functions in lthpruner library."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++# import g3++from absl.testing import parameterized+import numpy as np+import tensorflow as tf+print("Running TF", tf.__version__)++# TODO(b/139939526): move to public API.+from tensorflow.python.keras import keras_parameterized+from tensorflow_model_optimization.python.core.keras import compat+from tensorflow_model_optimization.python.core.sparsity.keras import pruning_schedule+from tensorflow_model_optimization.python.core.sparsity.keras import pruning_utils+from tensorflow_model_optimization.python.core.sparsity_tf2 import lthpruner as pruner++dtypes = tf.dtypes+test = tf.test++def get_lth_sparsity(save_round, n_rounds, target_sparsity, end_epoch):+  """+  save_round: when to save the weight; 1 onwards+  n_rounds: number of pruning cycles to do+  target_sparsity: the sparsity percentage to achieve by the end of the iteration+  end_epoch: total number of epochs to train for where pruning is eligible++  Returns:+    percent to prune to after each cyle/round+  """+  # no pruning until weights are saved TODO: off by one??+  n_rounds = tf.constant(n_rounds, dtype='float32') # dtype='float32'+  frequency = tf.math.floordiv(end_epoch - save_round + 1, n_rounds) # range(0, end, freq)+  prune_ratio_per_round = tf.math.pow(target_sparsity, tf.math.divide(1, n_rounds))+  return tf.cast(frequency, tf.int64), prune_ratio_per_round++def make_pruning_schedule(target_sparsity, begin, end, freq):+  return pruning_schedule.ConstantSparsity(target_sparsity, begin, end, freq)++def sample_noise(x, mu=0, sigma=1.):+  sample = tf.random.normal((), mean=mu,  stddev=sigma, dtype=tf.float64)+  return sample++def _dummy_gradient(x, dtype=tf.float32):+  try:+    base_type = x.dtype+  except:+    base_type = dtype+  grad = tf.ones_like(x, dtype=base_type)+  return grad++class PruningTest(test.TestCase, parameterized.TestCase):++  def setUp(self):+    super(PruningTest, self).setUp()+    self.block_size = (1, 1)+    self.block_pooling_type = "AVG"+    self.target_sparsity = 0.5+    self.constant_sparsity = pruning_schedule.ConstantSparsity(self.target_sparsity, 0, 100, 1)+    self.save_init = tf.Variable(0)+    self.save_itr_10 = 10+    self.dummy_optimizer = tf.keras.optimizers.SGD(learning_rate=0.01)+    self.grad = _dummy_gradient++  # Variable initialization outside of setUp() is needed for compatibility with+  # run_all_keras_modes.+  #+  # setUp() lies outside of the "eager scope" that wraps the test cases+  # themselves, resulting in initializing graph tensors instead of eager+  # tensors when testing eager execution.++  def testUpdateSingleMask(self):+    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")+    weight_dtype = weight.dtype.base_dtype+    mask = tf.Variable(+        tf.ones(weight.get_shape(), dtype=weight_dtype),+        name="mask",+        dtype=weight_dtype)+    threshold = tf.Variable(+        tf.zeros([], dtype=weight_dtype), name="threshold", dtype=weight_dtype)+    pruning_vars = [(weight, mask, threshold)]++    p = pruner.LTHPruner(+        pruning_schedule=self.constant_sparsity,+        save_iteration=self.save_init,+        block_size=self.block_size,+        block_pooling_type=self.block_pooling_type)++    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)+    p.create_slots(optimizer, weight)++    mask_before_pruning = optimizer.get_slot(weight, "mask").read_value()+    self.assertAllEqual(np.count_nonzero(mask_before_pruning), 100)++    next_step = optimizer.iterations.assign_add(1)+    p.update_masks(pruning_vars, next_step)++    mask_after_pruning = mask.read_value()+    self.assertAllEqual(np.count_nonzero(mask_after_pruning), 50)++  def testConstructsMaskAndThresholdCorrectly(self):+    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")+    weight_dtype = weight.dtype.base_dtype+    p = pruner.LTHPruner(+        # Sparsity math often returns values with small tolerances.+        pruning_schedule=lambda x: (True, 0.200000018),+        save_iteration=self.save_init,+        block_size=(1, 1), block_pooling_type=None)+        +    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)+    p.create_slots(optimizer, weight)+    step = optimizer.iterations++    # input matrix is [ 1.0, 2.0, ..., 8.0, 9.0, 10.0 ]+    threshold, mask = p._update_mask(step, np.arange(1, 11))++    self.assertEqual(3, threshold)+    self.assertAllEqual(+        # expected matrix is [ 0.0, 0.0, 1.0, 1.0 ... 1.0 ]+        np.concatenate((np.zeros(2), np.ones(8))), mask)++  def _blockMasking(self, block_size, block_pooling_type, weight,+                    expected_mask):+    mask = tf.Variable(+        tf.ones(weight.get_shape(), dtype=weight.dtype),+        name="mask",+        dtype=weight.dtype)+    threshold = tf.Variable(+        tf.zeros([], dtype=weight.dtype), name="threshold", dtype=weight.dtype)++    # Set up pruning+    p = pruner.LTHPruner(+        pruning_schedule=self.constant_sparsity,+        save_iteration=self.save_init,+        block_size=block_size,+        block_pooling_type=block_pooling_type)++    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)+    p.create_slots(optimizer, weight)+    step = optimizer.iterations++    _, new_mask = p._maybe_update_block_mask(step, weight)+    # Check if the mask is the same size as the weights+    self.assertAllEqual(new_mask.get_shape(), weight.get_shape())+    mask_after_pruning = new_mask+    self.assertAllEqual(mask_after_pruning, expected_mask)++  def testBlockMaskingAvg(self):+    block_size = (2, 2)+    block_pooling_type = "AVG"+    weight = tf.Variable([[0.1, 0.1, 0.2, 0.2], [0.1, 0.1, 0.2, 0.2],+                          [0.3, 0.3, 0.4, 0.4], [0.3, 0.3, 0.4, 0.4]])+    expected_mask = [[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0],+                     [1., 1., 1., 1.], [1., 1., 1., 1.]]++    self._blockMasking(block_size, block_pooling_type, weight, expected_mask)++  def testBlockMaskingMax(self):+    block_size = (2, 2)+    block_pooling_type = "MAX"+    weight = tf.Variable([[0.1, 0.0, 0.2, 0.0], [0.0, -0.1, 0.0, -0.2],+                                   [0.3, 0.0, 0.4, 0.0], [0.0, -0.3, 0.0,+                                                          -0.4]])+    expected_mask = [[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0],+                     [1., 1., 1., 1.], [1., 1., 1., 1.]]++    self._blockMasking(block_size, block_pooling_type, weight, expected_mask)++  def testBlockMaskingWithHigherDimensionsRaisesError(self):+    block_size = (2, 2)+    block_pooling_type = "AVG"+    # Weights as in testBlockMasking, but with one extra dimension.+    weight = tf.Variable([[[0.1, 0.1, 0.2, 0.2], [0.1, 0.1, 0.2, 0.2],+                                    [0.3, 0.3, 0.4, 0.4], [0.3, 0.3, 0.4,+                                                           0.4]]])+    expected_mask = [[[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0],+                      [1., 1., 1., 1.], [1., 1., 1., 1.]]]++    # Block masking should only be used with 2 Dimensional weights.+    with self.assertRaises(ValueError):+      self._blockMasking(block_size, block_pooling_type, weight, expected_mask)++  def testConditionalMaskUpdate(self):+    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")+    weight_dtype = weight.dtype.base_dtype+    mask = tf.Variable(+        tf.ones(weight.get_shape(), dtype=weight_dtype),+        name="mask",+        dtype=weight_dtype)+    threshold = tf.Variable(+        tf.zeros([], dtype=weight_dtype), name="threshold", dtype=weight_dtype)+    pruning_vars = [(weight, mask, threshold)]++    def linear_sparsity(step):+      sparsity_val = tf.convert_to_tensor(+          [0.0, 0.1, 0.1, 0.3, 0.3, 0.5, 0.5, 0.5, 0.5, 0.5])+      return tf.convert_to_tensor(True), sparsity_val[step]++    def weight_mask_op(pruning_vars):+      values_and_vars = []+      for weight, mask, _ in  pruning_vars:+        weight.assign(tf.math.multiply(weight, mask))++    # Set up pruning+    p = pruner.LTHPruner(+        pruning_schedule=linear_sparsity,+        save_iteration=self.save_init,+        block_size=self.block_size,+        block_pooling_type=self.block_pooling_type)++    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)+    p.create_slots(optimizer, weight)+    step = optimizer.iterations++    non_zero_count = []+    for _ in range(10):+      step = optimizer.iterations+      p.update_masks(pruning_vars, step)+      weight_mask_op(pruning_vars)+      optimizer.iterations.assign_add(1)++      non_zero_count.append(tf.math.count_nonzero(weight))++    # Weights pruned at steps 1,3,5+    expected_non_zero_count = [100, 90, 90, 70, 70, 50, 50, 50, 50, 50]+    self.assertAllEqual(expected_non_zero_count, non_zero_count)+++  def testSaveOriginalInitializations(self):+    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")+    weight_dtype = weight.dtype.base_dtype+    save_round = 0+    n_rounds = 5+    end_iter = 100+    frequency, prune_ratio_per_round = get_lth_sparsity(save_round, n_rounds, self.target_sparsity, end_iter)+    pruning_schedule = make_pruning_schedule(prune_ratio_per_round, 0, end_iter, frequency)+    +    p = pruner.LTHPruner(+        pruning_schedule=pruning_schedule,+        save_iteration=self.save_init,+        block_size=self.block_size,+        block_pooling_type=self.block_pooling_type)++    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)+    init_weights_before_pruning = weight.read_value()++    def _train():+      p.create_slots(optimizer, weight)+      p.preprocess_weights(optimizer, weight, self.grad(weight))+      weight.assign(tf.math.add(weight, sample_noise(0))) # update weights+      p.postprocess_weights(optimizer, weight, self.grad(weight))+      optimizer.iterations.assign_add(1)+      initialization_slot = optimizer.get_slot(weight, "original_initialization")+      return initialization_slot+    +    initialization_slot = _train()+    self.assertAllEqual(initialization_slot, init_weights_before_pruning)+    +    initialization_slot = tf.function(_train)()+    self.assertAllEqual(initialization_slot, init_weights_before_pruning)+++  def testSaveWeightsIterK(self):+    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")++    save_round = 5+    n_rounds = 24+    end_iter = 100+    frequency, prune_ratio_per_round = get_lth_sparsity(save_round, n_rounds, self.target_sparsity, end_iter)+    pruning_schedule = make_pruning_schedule(1 - prune_ratio_per_round, save_round, end_iter, frequency)+    +    p = pruner.LTHPruner(+        pruning_schedule=pruning_schedule,+        save_iteration=save_round,+        block_size=self.block_size,+        block_pooling_type=self.block_pooling_type)++    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)++    def _train(weight):+      expected_saved_initialization = tf.ones_like(weight, dtype=weight.dtype.base_dtype) * -1+      p.create_slots(optimizer, weight)+      for i in range(7):+        p.preprocess_weights(optimizer, weight, self.grad(weight))+        if optimizer.iterations == save_round:+          expected_saved_initialization = weight.read_value()+        weight.assign(tf.math.add(weight, sample_noise(i)))+        p.postprocess_weights(optimizer, weight, self.grad(weight))+        optimizer.iterations.assign_add(1)+      return expected_saved_initialization+      +    expected_saved_initialization = _train(weight)+    self.assertAllEqual(optimizer.get_slot(weight, "original_initialization"), expected_saved_initialization)++    initialization_slot_k = tf.math.multiply(optimizer.get_slot(weight, "original_initialization"), optimizer.get_slot(weight, "mask"))+    masked_expected = tf.math.multiply(expected_saved_initialization, optimizer.get_slot(weight, "mask"))+    self.assertAllEqual(initialization_slot_k, masked_expected)++    mask_after_pruning = optimizer.get_slot(weight, "mask").read_value()+    self.assertAllEqual(np.count_nonzero(mask_after_pruning), 97)++    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")+    optimizer.iterations.assign(0)+    expected_saved_initialization = tf.function(_train)(weight)+    self.assertAllEqual(optimizer.get_slot(weight, "original_initialization"), expected_saved_initialization)++    initialization_slot_k = tf.math.multiply(optimizer.get_slot(weight, "original_initialization"), optimizer.get_slot(weight, "mask"))+    masked_expected = tf.math.multiply(expected_saved_initialization, optimizer.get_slot(weight, "mask"))+    self.assertAllEqual(initialization_slot_k, masked_expected)++    mask_after_pruning = optimizer.get_slot(weight, "mask").read_value()+    self.assertAllEqual(np.count_nonzero(mask_after_pruning), 97)++  def testReloadWeightsatInitialization(self):+    weight1 = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")++    save_round = 0+    n_rounds = 20+    end_iter = 100+    frequency, prune_ratio_per_round = get_lth_sparsity(save_round, n_rounds, self.target_sparsity, end_iter)+    pruning_schedule = make_pruning_schedule(1 - prune_ratio_per_round, save_round, end_iter, frequency)+    +    p = pruner.LTHPruner(+        pruning_schedule=pruning_schedule,+        save_iteration=self.save_init,+        block_size=self.block_size,+        block_pooling_type=self.block_pooling_type)++    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)+    expected_saved_initialization = None++    def _train(weight):+      expected_saved_initialization = tf.ones_like(weight, dtype=weight.dtype.base_dtype) * -1+      p.create_slots(optimizer, weight)+      for i in tf.range(0, 2):+        p.preprocess_weights(optimizer, weight, self.grad(weight))+        if i == save_round:+          expected_saved_initialization = weight.read_value()+        weight.assign(tf.math.add(weight, sample_noise(i)))+        p.postprocess_weights(optimizer, weight, self.grad(weight))+        optimizer.iterations.assign_add(1) # save weights right before iteration+      return expected_saved_initialization++    expected_saved_initialization = _train(weight1)+    initialization_slot = optimizer.get_slot(weight1, "original_initialization")+    self.assertAllEqual(initialization_slot, expected_saved_initialization)++    mask_after_pruning = optimizer.get_slot(weight1, "mask").read_value()+    masked_weight_expected = tf.math.multiply(initialization_slot, mask_after_pruning)+    self.assertAllEqual(np.count_nonzero(masked_weight_expected), 97)+    self.assertAllEqual(np.count_nonzero(mask_after_pruning), 97)++    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")+    optimizer.iterations.assign(0)+    expected_saved_initialization = tf.function(_train)(weight)+    initialization_slot = optimizer.get_slot(weight, "original_initialization")+    self.assertAllEqual(initialization_slot, expected_saved_initialization)++    mask_after_pruning = optimizer.get_slot(weight, "mask").read_value()+    masked_weight_expected = tf.math.multiply(initialization_slot, mask_after_pruning)+    self.assertAllEqual(np.count_nonzero(masked_weight_expected), 97)+    self.assertAllEqual(np.count_nonzero(mask_after_pruning), 97)++  +  def testReloadWeightsIterationK(self):+    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")++    save_round = 5+    n_rounds = 24+    end_iter = 100+    frequency, prune_ratio_per_round = get_lth_sparsity(save_round, n_rounds, self.target_sparsity, end_iter)+    pruning_schedule = make_pruning_schedule(1 - prune_ratio_per_round, save_round + 1, end_iter, frequency)+    +    p = pruner.LTHPruner(+        pruning_schedule=pruning_schedule,+        save_iteration=save_round,+        block_size=self.block_size,+        block_pooling_type=self.block_pooling_type)++    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)++    def _train(weight):+      p.create_slots(optimizer, weight)+      for i in range(7):+        p.preprocess_weights(optimizer, weight, self.grad(weight))+        weight.assign(tf.math.add(weight, sample_noise(i)))+        p.postprocess_weights(optimizer, weight, self.grad(weight))+        optimizer.iterations.assign_add(1)++    _train(weight)+    masked_orig_init = tf.math.multiply(optimizer.get_slot(weight, "original_initialization"), optimizer.get_slot(weight, 'mask'))+    self.assertAllEqual(masked_orig_init, weight)++    mask_after_pruning = optimizer.get_slot(weight, "mask").read_value()+    self.assertAllEqual(np.count_nonzero(mask_after_pruning), 97)++    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")+    optimizer.iterations.assign(0)+    tf.function(_train)(weight)+    masked_orig_init = tf.math.multiply(optimizer.get_slot(weight, "original_initialization"), optimizer.get_slot(weight, 'mask'))+    self.assertAllEqual(masked_orig_init, weight)++    mask_after_pruning = optimizer.get_slot(weight, "mask").read_value()+    self.assertAllEqual(np.count_nonzero(mask_after_pruning), 97)+++  def testReloadTwoTimes(self):+    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")+    weight_dtype = weight.dtype.base_dtype++    save_round = 5+    n_rounds = 24+    end_iter = 100+    frequency, prune_ratio_per_round = get_lth_sparsity(save_round, n_rounds, self.target_sparsity, end_iter)+    pruning_schedule = make_pruning_schedule(1 - prune_ratio_per_round, save_round + 1, end_iter, frequency)+    +    p = pruner.LTHPruner(+        pruning_schedule=pruning_schedule,+        save_iteration=save_round,+        block_size=self.block_size,+        block_pooling_type=self.block_pooling_type)++    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)++    def _train(weight):+      expected_first_saved_initialization = None+      expected_second_saved_initialization = None+      p.create_slots(optimizer, weight)+      for i in range(0, 13): # this should save, reload, and update exactly twice+        p.preprocess_weights(optimizer, weight, self.grad(weight))+        if i == save_round - 1:+          expected_first_saved_initialization = optimizer.get_slot(weight, "original_initialization")+        if i == (save_round - 1) * 2:+          expected_second_saved_initialization = optimizer.get_slot(weight, "original_initialization")+        weight.assign(tf.math.add(weight, sample_noise(i)))+        p.postprocess_weights(optimizer, weight, self.grad(weight))+        optimizer.iterations.assign_add(1)+      return expected_first_saved_initialization, expected_second_saved_initialization++    expected_first_saved_initialization, expected_second_saved_initialization = _train(weight)++    self.assertAllEqual(expected_first_saved_initialization, expected_second_saved_initialization)++    initialization_slot = optimizer.get_slot(weight, "original_initialization")+    self.assertAllEqual(initialization_slot, expected_first_saved_initialization)

+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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 key functions in lthpruner library."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++# import g3++from absl.testing import parameterized+import numpy as np+import tensorflow as tf+print("Running TF", tf.__version__)++# TODO(b/139939526): move to public API.+from tensorflow.python.keras import keras_parameterized+from tensorflow_model_optimization.python.core.keras import compat+from tensorflow_model_optimization.python.core.sparsity.keras import pruning_schedule+from tensorflow_model_optimization.python.core.sparsity.keras import pruning_utils+from tensorflow_model_optimization.python.core.sparsity_tf2 import lthpruner as pruner++dtypes = tf.dtypes+test = tf.test++def get_lth_sparsity(save_round, n_rounds, target_sparsity, end_epoch):+  """+  save_round: when to save the weight; 1 onwards+  n_rounds: number of pruning cycles to do+  target_sparsity: the sparsity percentage to achieve by the end of the iteration+  end_epoch: total number of epochs to train for where pruning is eligible++  Returns:+    percent to prune to after each cyle/round+  """+  # no pruning until weights are saved TODO: off by one??+  n_rounds = tf.constant(n_rounds, dtype='float32') # dtype='float32'+  frequency = tf.math.floordiv(end_epoch - save_round + 1, n_rounds) # range(0, end, freq)+  prune_ratio_per_round = tf.math.pow(target_sparsity, tf.math.divide(1, n_rounds))+  return tf.cast(frequency, tf.int64), prune_ratio_per_round++def make_pruning_schedule(target_sparsity, begin, end, freq):+  return pruning_schedule.ConstantSparsity(target_sparsity, begin, end, freq)++def sample_noise(x, mu=0, sigma=1.):+  sample = tf.random.normal((), mean=mu,  stddev=sigma, dtype=tf.float64)+  return sample++def _dummy_gradient(x, dtype=tf.float32):+  try:+    base_type = x.dtype+  except:+    base_type = dtype+  grad = tf.ones_like(x, dtype=base_type)+  return grad++class PruningTest(test.TestCase, parameterized.TestCase):++  def setUp(self):+    super(PruningTest, self).setUp()+    self.block_size = (1, 1)+    self.block_pooling_type = "AVG"+    self.target_sparsity = 0.5+    self.constant_sparsity = pruning_schedule.ConstantSparsity(self.target_sparsity, 0, 100, 1)+    self.save_init = tf.Variable(0)+    self.save_itr_10 = 10+    self.dummy_optimizer = tf.keras.optimizers.SGD(learning_rate=0.01)+    self.grad = _dummy_gradient++  # Variable initialization outside of setUp() is needed for compatibility with+  # run_all_keras_modes.+  #+  # setUp() lies outside of the "eager scope" that wraps the test cases+  # themselves, resulting in initializing graph tensors instead of eager+  # tensors when testing eager execution.++  def testUpdateSingleMask(self):+    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")+    weight_dtype = weight.dtype.base_dtype+    mask = tf.Variable(+        tf.ones(weight.get_shape(), dtype=weight_dtype),+        name="mask",+        dtype=weight_dtype)+    threshold = tf.Variable(+        tf.zeros([], dtype=weight_dtype), name="threshold", dtype=weight_dtype)+    pruning_vars = [(weight, mask, threshold)]++    p = pruner.LTHPruner(+        pruning_schedule=self.constant_sparsity,+        save_iteration=self.save_init,+        block_size=self.block_size,+        block_pooling_type=self.block_pooling_type)++    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)+    p.create_slots(optimizer, weight)++    mask_before_pruning = optimizer.get_slot(weight, "mask").read_value()+    self.assertAllEqual(np.count_nonzero(mask_before_pruning), 100)++    next_step = optimizer.iterations.assign_add(1)+    p.update_masks(pruning_vars, next_step)++    mask_after_pruning = mask.read_value()+    self.assertAllEqual(np.count_nonzero(mask_after_pruning), 50)++  def testConstructsMaskAndThresholdCorrectly(self):+    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")+    weight_dtype = weight.dtype.base_dtype+    p = pruner.LTHPruner(+        # Sparsity math often returns values with small tolerances.+        pruning_schedule=lambda x: (True, 0.200000018),+        save_iteration=self.save_init,+        block_size=(1, 1), block_pooling_type=None)+        +    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)+    p.create_slots(optimizer, weight)+    step = optimizer.iterations++    # input matrix is [ 1.0, 2.0, ..., 8.0, 9.0, 10.0 ]+    threshold, mask = p._update_mask(step, np.arange(1, 11))++    self.assertEqual(3, threshold)+    self.assertAllEqual(+        # expected matrix is [ 0.0, 0.0, 1.0, 1.0 ... 1.0 ]+        np.concatenate((np.zeros(2), np.ones(8))), mask)++  def _blockMasking(self, block_size, block_pooling_type, weight,+                    expected_mask):+    mask = tf.Variable(+        tf.ones(weight.get_shape(), dtype=weight.dtype),+        name="mask",+        dtype=weight.dtype)+    threshold = tf.Variable(+        tf.zeros([], dtype=weight.dtype), name="threshold", dtype=weight.dtype)++    # Set up pruning+    p = pruner.LTHPruner(+        pruning_schedule=self.constant_sparsity,+        save_iteration=self.save_init,+        block_size=block_size,+        block_pooling_type=block_pooling_type)++    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)+    p.create_slots(optimizer, weight)+    step = optimizer.iterations++    _, new_mask = p._maybe_update_block_mask(step, weight)+    # Check if the mask is the same size as the weights+    self.assertAllEqual(new_mask.get_shape(), weight.get_shape())+    mask_after_pruning = new_mask+    self.assertAllEqual(mask_after_pruning, expected_mask)++  def testBlockMaskingAvg(self):+    block_size = (2, 2)+    block_pooling_type = "AVG"+    weight = tf.Variable([[0.1, 0.1, 0.2, 0.2], [0.1, 0.1, 0.2, 0.2],+                          [0.3, 0.3, 0.4, 0.4], [0.3, 0.3, 0.4, 0.4]])+    expected_mask = [[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0],+                     [1., 1., 1., 1.], [1., 1., 1., 1.]]++    self._blockMasking(block_size, block_pooling_type, weight, expected_mask)++  def testBlockMaskingMax(self):+    block_size = (2, 2)+    block_pooling_type = "MAX"+    weight = tf.Variable([[0.1, 0.0, 0.2, 0.0], [0.0, -0.1, 0.0, -0.2],+                                   [0.3, 0.0, 0.4, 0.0], [0.0, -0.3, 0.0,+                                                          -0.4]])+    expected_mask = [[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0],+                     [1., 1., 1., 1.], [1., 1., 1., 1.]]++    self._blockMasking(block_size, block_pooling_type, weight, expected_mask)++  def testBlockMaskingWithHigherDimensionsRaisesError(self):+    block_size = (2, 2)+    block_pooling_type = "AVG"+    # Weights as in testBlockMasking, but with one extra dimension.+    weight = tf.Variable([[[0.1, 0.1, 0.2, 0.2], [0.1, 0.1, 0.2, 0.2],+                                    [0.3, 0.3, 0.4, 0.4], [0.3, 0.3, 0.4,+                                                           0.4]]])+    expected_mask = [[[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0],+                      [1., 1., 1., 1.], [1., 1., 1., 1.]]]++    # Block masking should only be used with 2 Dimensional weights.+    with self.assertRaises(ValueError):+      self._blockMasking(block_size, block_pooling_type, weight, expected_mask)++  def testConditionalMaskUpdate(self):+    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")+    weight_dtype = weight.dtype.base_dtype+    mask = tf.Variable(+        tf.ones(weight.get_shape(), dtype=weight_dtype),+        name="mask",+        dtype=weight_dtype)+    threshold = tf.Variable(+        tf.zeros([], dtype=weight_dtype), name="threshold", dtype=weight_dtype)+    pruning_vars = [(weight, mask, threshold)]++    def linear_sparsity(step):+      sparsity_val = tf.convert_to_tensor(+          [0.0, 0.1, 0.1, 0.3, 0.3, 0.5, 0.5, 0.5, 0.5, 0.5])+      return tf.convert_to_tensor(True), sparsity_val[step]++    def weight_mask_op(pruning_vars):+      values_and_vars = []+      for weight, mask, _ in  pruning_vars:+        weight.assign(tf.math.multiply(weight, mask))++    # Set up pruning+    p = pruner.LTHPruner(+        pruning_schedule=linear_sparsity,+        save_iteration=self.save_init,+        block_size=self.block_size,+        block_pooling_type=self.block_pooling_type)++    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)+    p.create_slots(optimizer, weight)+    step = optimizer.iterations++    non_zero_count = []+    for _ in range(10):+      step = optimizer.iterations+      p.update_masks(pruning_vars, step)+      weight_mask_op(pruning_vars)+      optimizer.iterations.assign_add(1)++      non_zero_count.append(tf.math.count_nonzero(weight))++    # Weights pruned at steps 1,3,5+    expected_non_zero_count = [100, 90, 90, 70, 70, 50, 50, 50, 50, 50]+    self.assertAllEqual(expected_non_zero_count, non_zero_count)+++  def testSaveOriginalInitializations(self):+    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")+    weight_dtype = weight.dtype.base_dtype+    save_round = 0+    n_rounds = 5+    end_iter = 100+    frequency, prune_ratio_per_round = get_lth_sparsity(save_round, n_rounds, self.target_sparsity, end_iter)+    pruning_schedule = make_pruning_schedule(prune_ratio_per_round, 0, end_iter, frequency)+    +    p = pruner.LTHPruner(+        pruning_schedule=pruning_schedule,+        save_iteration=self.save_init,+        block_size=self.block_size,+        block_pooling_type=self.block_pooling_type)++    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)+    init_weights_before_pruning = weight.read_value()++    def _train():+      p.create_slots(optimizer, weight)+      p.preprocess_weights(optimizer, weight, self.grad(weight))+      weight.assign(tf.math.add(weight, sample_noise(0))) # update weights+      p.postprocess_weights(optimizer, weight, self.grad(weight))+      optimizer.iterations.assign_add(1)+      initialization_slot = optimizer.get_slot(weight, "original_initialization")+      return initialization_slot+    +    initialization_slot = _train()+    self.assertAllEqual(initialization_slot, init_weights_before_pruning)+    +    initialization_slot = tf.function(_train)()+    self.assertAllEqual(initialization_slot, init_weights_before_pruning)+++  def testSaveWeightsIterK(self):+    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")++    save_round = 5+    n_rounds = 24+    end_iter = 100+    frequency, prune_ratio_per_round = get_lth_sparsity(save_round, n_rounds, self.target_sparsity, end_iter)+    pruning_schedule = make_pruning_schedule(1 - prune_ratio_per_round, save_round, end_iter, frequency)+    +    p = pruner.LTHPruner(+        pruning_schedule=pruning_schedule,+        save_iteration=save_round,+        block_size=self.block_size,+        block_pooling_type=self.block_pooling_type)++    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)++    def _train(weight):+      expected_saved_initialization = tf.ones_like(weight, dtype=weight.dtype.base_dtype) * -1+      p.create_slots(optimizer, weight)+      for i in range(7):+        p.preprocess_weights(optimizer, weight, self.grad(weight))+        if optimizer.iterations == save_round:+          expected_saved_initialization = weight.read_value()+        weight.assign(tf.math.add(weight, sample_noise(i)))+        p.postprocess_weights(optimizer, weight, self.grad(weight))+        optimizer.iterations.assign_add(1)+      return expected_saved_initialization+      +    expected_saved_initialization = _train(weight)+    self.assertAllEqual(optimizer.get_slot(weight, "original_initialization"), expected_saved_initialization)++    initialization_slot_k = tf.math.multiply(optimizer.get_slot(weight, "original_initialization"), optimizer.get_slot(weight, "mask"))+    masked_expected = tf.math.multiply(expected_saved_initialization, optimizer.get_slot(weight, "mask"))+    self.assertAllEqual(initialization_slot_k, masked_expected)++    mask_after_pruning = optimizer.get_slot(weight, "mask").read_value()+    self.assertAllEqual(np.count_nonzero(mask_after_pruning), 97)++    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")+    optimizer.iterations.assign(0)+    expected_saved_initialization = tf.function(_train)(weight)+    self.assertAllEqual(optimizer.get_slot(weight, "original_initialization"), expected_saved_initialization)++    initialization_slot_k = tf.math.multiply(optimizer.get_slot(weight, "original_initialization"), optimizer.get_slot(weight, "mask"))+    masked_expected = tf.math.multiply(expected_saved_initialization, optimizer.get_slot(weight, "mask"))+    self.assertAllEqual(initialization_slot_k, masked_expected)++    mask_after_pruning = optimizer.get_slot(weight, "mask").read_value()+    self.assertAllEqual(np.count_nonzero(mask_after_pruning), 97)++  def testReloadWeightsatInitialization(self):+    weight1 = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")++    save_round = 0+    n_rounds = 20+    end_iter = 100+    frequency, prune_ratio_per_round = get_lth_sparsity(save_round, n_rounds, self.target_sparsity, end_iter)+    pruning_schedule = make_pruning_schedule(1 - prune_ratio_per_round, save_round, end_iter, frequency)+    +    p = pruner.LTHPruner(+        pruning_schedule=pruning_schedule,+        save_iteration=self.save_init,+        block_size=self.block_size,+        block_pooling_type=self.block_pooling_type)++    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)+    expected_saved_initialization = None++    def _train(weight):+      expected_saved_initialization = tf.ones_like(weight, dtype=weight.dtype.base_dtype) * -1+      p.create_slots(optimizer, weight)+      for i in tf.range(0, 2):+        p.preprocess_weights(optimizer, weight, self.grad(weight))+        if i == save_round:+          expected_saved_initialization = weight.read_value()+        weight.assign(tf.math.add(weight, sample_noise(i)))+        p.postprocess_weights(optimizer, weight, self.grad(weight))+        optimizer.iterations.assign_add(1) # save weights right before iteration+      return expected_saved_initialization++    expected_saved_initialization = _train(weight1)+    initialization_slot = optimizer.get_slot(weight1, "original_initialization")+    self.assertAllEqual(initialization_slot, expected_saved_initialization)++    mask_after_pruning = optimizer.get_slot(weight1, "mask").read_value()+    masked_weight_expected = tf.math.multiply(initialization_slot, mask_after_pruning)+    self.assertAllEqual(np.count_nonzero(masked_weight_expected), 97)+    self.assertAllEqual(np.count_nonzero(mask_after_pruning), 97)++    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")+    optimizer.iterations.assign(0)+    expected_saved_initialization = tf.function(_train)(weight)+    initialization_slot = optimizer.get_slot(weight, "original_initialization")+    self.assertAllEqual(initialization_slot, expected_saved_initialization)++    mask_after_pruning = optimizer.get_slot(weight, "mask").read_value()+    masked_weight_expected = tf.math.multiply(initialization_slot, mask_after_pruning)+    self.assertAllEqual(np.count_nonzero(masked_weight_expected), 97)+    self.assertAllEqual(np.count_nonzero(mask_after_pruning), 97)++  +  def testReloadWeightsIterationK(self):+    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")++    save_round = 5+    n_rounds = 24+    end_iter = 100+    frequency, prune_ratio_per_round = get_lth_sparsity(save_round, n_rounds, self.target_sparsity, end_iter)+    pruning_schedule = make_pruning_schedule(1 - prune_ratio_per_round, save_round + 1, end_iter, frequency)+    +    p = pruner.LTHPruner(+        pruning_schedule=pruning_schedule,+        save_iteration=save_round,+        block_size=self.block_size,+        block_pooling_type=self.block_pooling_type)++    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)++    def _train(weight):+      p.create_slots(optimizer, weight)+      for i in range(7):+        p.preprocess_weights(optimizer, weight, self.grad(weight))+        weight.assign(tf.math.add(weight, sample_noise(i)))+        p.postprocess_weights(optimizer, weight, self.grad(weight))+        optimizer.iterations.assign_add(1)++    _train(weight)+    masked_orig_init = tf.math.multiply(optimizer.get_slot(weight, "original_initialization"), optimizer.get_slot(weight, 'mask'))+    self.assertAllEqual(masked_orig_init, weight)++    mask_after_pruning = optimizer.get_slot(weight, "mask").read_value()+    self.assertAllEqual(np.count_nonzero(mask_after_pruning), 97)++    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")+    optimizer.iterations.assign(0)+    tf.function(_train)(weight)+    masked_orig_init = tf.math.multiply(optimizer.get_slot(weight, "original_initialization"), optimizer.get_slot(weight, 'mask'))+    self.assertAllEqual(masked_orig_init, weight)++    mask_after_pruning = optimizer.get_slot(weight, "mask").read_value()+    self.assertAllEqual(np.count_nonzero(mask_after_pruning), 97)+++  def testReloadTwoTimes(self):+    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")+    weight_dtype = weight.dtype.base_dtype++    save_round = 5+    n_rounds = 24+    end_iter = 100+    frequency, prune_ratio_per_round = get_lth_sparsity(save_round, n_rounds, self.target_sparsity, end_iter)+    pruning_schedule = make_pruning_schedule(1 - prune_ratio_per_round, save_round + 1, end_iter, frequency)+    +    p = pruner.LTHPruner(+        pruning_schedule=pruning_schedule,+        save_iteration=save_round,+        block_size=self.block_size,+        block_pooling_type=self.block_pooling_type)++    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)++    def _train(weight):+      expected_first_saved_initialization = None+      expected_second_saved_initialization = None+      p.create_slots(optimizer, weight)+      for i in range(0, 13): # this should save, reload, and update exactly twice+        p.preprocess_weights(optimizer, weight, self.grad(weight))+        if i == save_round - 1:

+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.+#+# Licensed under the Apache License, Version 2.0 (the "License");+# you may not use this file except in compliance with the License.+# You may obtain a copy of the License at+#+#     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 key functions in lthpruner library."""++from __future__ import absolute_import+from __future__ import division+from __future__ import print_function++# import g3++from absl.testing import parameterized+import numpy as np+import tensorflow as tf+print("Running TF", tf.__version__)++# TODO(b/139939526): move to public API.+from tensorflow.python.keras import keras_parameterized+from tensorflow_model_optimization.python.core.keras import compat+from tensorflow_model_optimization.python.core.sparsity.keras import pruning_schedule+from tensorflow_model_optimization.python.core.sparsity.keras import pruning_utils+from tensorflow_model_optimization.python.core.sparsity_tf2 import lthpruner as pruner++dtypes = tf.dtypes+test = tf.test++def get_lth_sparsity(save_round, n_rounds, target_sparsity, end_epoch):+  """+  save_round: when to save the weight; 1 onwards+  n_rounds: number of pruning cycles to do+  target_sparsity: the sparsity percentage to achieve by the end of the iteration+  end_epoch: total number of epochs to train for where pruning is eligible++  Returns:+    percent to prune to after each cyle/round+  """+  # no pruning until weights are saved TODO: off by one??+  n_rounds = tf.constant(n_rounds, dtype='float32') # dtype='float32'+  frequency = tf.math.floordiv(end_epoch - save_round + 1, n_rounds) # range(0, end, freq)+  prune_ratio_per_round = tf.math.pow(target_sparsity, tf.math.divide(1, n_rounds))+  return tf.cast(frequency, tf.int64), prune_ratio_per_round++def make_pruning_schedule(target_sparsity, begin, end, freq):+  return pruning_schedule.ConstantSparsity(target_sparsity, begin, end, freq)++def sample_noise(x, mu=0, sigma=1.):+  sample = tf.random.normal((), mean=mu,  stddev=sigma, dtype=tf.float64)+  return sample++def _dummy_gradient(x, dtype=tf.float32):+  try:+    base_type = x.dtype+  except:+    base_type = dtype+  grad = tf.ones_like(x, dtype=base_type)+  return grad++class PruningTest(test.TestCase, parameterized.TestCase):++  def setUp(self):+    super(PruningTest, self).setUp()+    self.block_size = (1, 1)+    self.block_pooling_type = "AVG"+    self.target_sparsity = 0.5+    self.constant_sparsity = pruning_schedule.ConstantSparsity(self.target_sparsity, 0, 100, 1)+    self.save_init = tf.Variable(0)+    self.save_itr_10 = 10+    self.dummy_optimizer = tf.keras.optimizers.SGD(learning_rate=0.01)+    self.grad = _dummy_gradient++  # Variable initialization outside of setUp() is needed for compatibility with+  # run_all_keras_modes.+  #+  # setUp() lies outside of the "eager scope" that wraps the test cases+  # themselves, resulting in initializing graph tensors instead of eager+  # tensors when testing eager execution.++  def testUpdateSingleMask(self):+    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")+    weight_dtype = weight.dtype.base_dtype+    mask = tf.Variable(+        tf.ones(weight.get_shape(), dtype=weight_dtype),+        name="mask",+        dtype=weight_dtype)+    threshold = tf.Variable(+        tf.zeros([], dtype=weight_dtype), name="threshold", dtype=weight_dtype)+    pruning_vars = [(weight, mask, threshold)]++    p = pruner.LTHPruner(+        pruning_schedule=self.constant_sparsity,+        save_iteration=self.save_init,+        block_size=self.block_size,+        block_pooling_type=self.block_pooling_type)++    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)+    p.create_slots(optimizer, weight)++    mask_before_pruning = optimizer.get_slot(weight, "mask").read_value()+    self.assertAllEqual(np.count_nonzero(mask_before_pruning), 100)++    next_step = optimizer.iterations.assign_add(1)+    p.update_masks(pruning_vars, next_step)++    mask_after_pruning = mask.read_value()+    self.assertAllEqual(np.count_nonzero(mask_after_pruning), 50)++  def testConstructsMaskAndThresholdCorrectly(self):+    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")+    weight_dtype = weight.dtype.base_dtype+    p = pruner.LTHPruner(+        # Sparsity math often returns values with small tolerances.+        pruning_schedule=lambda x: (True, 0.200000018),+        save_iteration=self.save_init,+        block_size=(1, 1), block_pooling_type=None)+        +    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)+    p.create_slots(optimizer, weight)+    step = optimizer.iterations++    # input matrix is [ 1.0, 2.0, ..., 8.0, 9.0, 10.0 ]+    threshold, mask = p._update_mask(step, np.arange(1, 11))++    self.assertEqual(3, threshold)+    self.assertAllEqual(+        # expected matrix is [ 0.0, 0.0, 1.0, 1.0 ... 1.0 ]+        np.concatenate((np.zeros(2), np.ones(8))), mask)++  def _blockMasking(self, block_size, block_pooling_type, weight,+                    expected_mask):+    mask = tf.Variable(+        tf.ones(weight.get_shape(), dtype=weight.dtype),+        name="mask",+        dtype=weight.dtype)+    threshold = tf.Variable(+        tf.zeros([], dtype=weight.dtype), name="threshold", dtype=weight.dtype)++    # Set up pruning+    p = pruner.LTHPruner(+        pruning_schedule=self.constant_sparsity,+        save_iteration=self.save_init,+        block_size=block_size,+        block_pooling_type=block_pooling_type)++    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)+    p.create_slots(optimizer, weight)+    step = optimizer.iterations++    _, new_mask = p._maybe_update_block_mask(step, weight)+    # Check if the mask is the same size as the weights+    self.assertAllEqual(new_mask.get_shape(), weight.get_shape())+    mask_after_pruning = new_mask+    self.assertAllEqual(mask_after_pruning, expected_mask)++  def testBlockMaskingAvg(self):+    block_size = (2, 2)+    block_pooling_type = "AVG"+    weight = tf.Variable([[0.1, 0.1, 0.2, 0.2], [0.1, 0.1, 0.2, 0.2],+                          [0.3, 0.3, 0.4, 0.4], [0.3, 0.3, 0.4, 0.4]])+    expected_mask = [[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0],+                     [1., 1., 1., 1.], [1., 1., 1., 1.]]++    self._blockMasking(block_size, block_pooling_type, weight, expected_mask)++  def testBlockMaskingMax(self):+    block_size = (2, 2)+    block_pooling_type = "MAX"+    weight = tf.Variable([[0.1, 0.0, 0.2, 0.0], [0.0, -0.1, 0.0, -0.2],+                                   [0.3, 0.0, 0.4, 0.0], [0.0, -0.3, 0.0,+                                                          -0.4]])+    expected_mask = [[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0],+                     [1., 1., 1., 1.], [1., 1., 1., 1.]]++    self._blockMasking(block_size, block_pooling_type, weight, expected_mask)++  def testBlockMaskingWithHigherDimensionsRaisesError(self):+    block_size = (2, 2)+    block_pooling_type = "AVG"+    # Weights as in testBlockMasking, but with one extra dimension.+    weight = tf.Variable([[[0.1, 0.1, 0.2, 0.2], [0.1, 0.1, 0.2, 0.2],+                                    [0.3, 0.3, 0.4, 0.4], [0.3, 0.3, 0.4,+                                                           0.4]]])+    expected_mask = [[[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0],+                      [1., 1., 1., 1.], [1., 1., 1., 1.]]]++    # Block masking should only be used with 2 Dimensional weights.+    with self.assertRaises(ValueError):+      self._blockMasking(block_size, block_pooling_type, weight, expected_mask)++  def testConditionalMaskUpdate(self):+    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")+    weight_dtype = weight.dtype.base_dtype+    mask = tf.Variable(+        tf.ones(weight.get_shape(), dtype=weight_dtype),+        name="mask",+        dtype=weight_dtype)+    threshold = tf.Variable(+        tf.zeros([], dtype=weight_dtype), name="threshold", dtype=weight_dtype)+    pruning_vars = [(weight, mask, threshold)]++    def linear_sparsity(step):+      sparsity_val = tf.convert_to_tensor(+          [0.0, 0.1, 0.1, 0.3, 0.3, 0.5, 0.5, 0.5, 0.5, 0.5])+      return tf.convert_to_tensor(True), sparsity_val[step]++    def weight_mask_op(pruning_vars):+      values_and_vars = []+      for weight, mask, _ in  pruning_vars:+        weight.assign(tf.math.multiply(weight, mask))++    # Set up pruning+    p = pruner.LTHPruner(+        pruning_schedule=linear_sparsity,+        save_iteration=self.save_init,+        block_size=self.block_size,+        block_pooling_type=self.block_pooling_type)++    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)+    p.create_slots(optimizer, weight)+    step = optimizer.iterations++    non_zero_count = []+    for _ in range(10):+      step = optimizer.iterations+      p.update_masks(pruning_vars, step)+      weight_mask_op(pruning_vars)+      optimizer.iterations.assign_add(1)++      non_zero_count.append(tf.math.count_nonzero(weight))++    # Weights pruned at steps 1,3,5+    expected_non_zero_count = [100, 90, 90, 70, 70, 50, 50, 50, 50, 50]+    self.assertAllEqual(expected_non_zero_count, non_zero_count)+++  def testSaveOriginalInitializations(self):+    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")+    weight_dtype = weight.dtype.base_dtype+    save_round = 0+    n_rounds = 5+    end_iter = 100+    frequency, prune_ratio_per_round = get_lth_sparsity(save_round, n_rounds, self.target_sparsity, end_iter)+    pruning_schedule = make_pruning_schedule(prune_ratio_per_round, 0, end_iter, frequency)+    +    p = pruner.LTHPruner(+        pruning_schedule=pruning_schedule,+        save_iteration=self.save_init,+        block_size=self.block_size,+        block_pooling_type=self.block_pooling_type)++    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)+    init_weights_before_pruning = weight.read_value()++    def _train():+      p.create_slots(optimizer, weight)+      p.preprocess_weights(optimizer, weight, self.grad(weight))+      weight.assign(tf.math.add(weight, sample_noise(0))) # update weights+      p.postprocess_weights(optimizer, weight, self.grad(weight))+      optimizer.iterations.assign_add(1)+      initialization_slot = optimizer.get_slot(weight, "original_initialization")+      return initialization_slot+    +    initialization_slot = _train()+    self.assertAllEqual(initialization_slot, init_weights_before_pruning)+    +    initialization_slot = tf.function(_train)()+    self.assertAllEqual(initialization_slot, init_weights_before_pruning)+++  def testSaveWeightsIterK(self):+    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")++    save_round = 5+    n_rounds = 24+    end_iter = 100+    frequency, prune_ratio_per_round = get_lth_sparsity(save_round, n_rounds, self.target_sparsity, end_iter)+    pruning_schedule = make_pruning_schedule(1 - prune_ratio_per_round, save_round, end_iter, frequency)+    +    p = pruner.LTHPruner(+        pruning_schedule=pruning_schedule,+        save_iteration=save_round,+        block_size=self.block_size,+        block_pooling_type=self.block_pooling_type)++    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)++    def _train(weight):+      expected_saved_initialization = tf.ones_like(weight, dtype=weight.dtype.base_dtype) * -1+      p.create_slots(optimizer, weight)+      for i in range(7):+        p.preprocess_weights(optimizer, weight, self.grad(weight))+        if optimizer.iterations == save_round:+          expected_saved_initialization = weight.read_value()+        weight.assign(tf.math.add(weight, sample_noise(i)))+        p.postprocess_weights(optimizer, weight, self.grad(weight))+        optimizer.iterations.assign_add(1)+      return expected_saved_initialization+      +    expected_saved_initialization = _train(weight)+    self.assertAllEqual(optimizer.get_slot(weight, "original_initialization"), expected_saved_initialization)++    initialization_slot_k = tf.math.multiply(optimizer.get_slot(weight, "original_initialization"), optimizer.get_slot(weight, "mask"))+    masked_expected = tf.math.multiply(expected_saved_initialization, optimizer.get_slot(weight, "mask"))+    self.assertAllEqual(initialization_slot_k, masked_expected)++    mask_after_pruning = optimizer.get_slot(weight, "mask").read_value()+    self.assertAllEqual(np.count_nonzero(mask_after_pruning), 97)++    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")+    optimizer.iterations.assign(0)+    expected_saved_initialization = tf.function(_train)(weight)+    self.assertAllEqual(optimizer.get_slot(weight, "original_initialization"), expected_saved_initialization)++    initialization_slot_k = tf.math.multiply(optimizer.get_slot(weight, "original_initialization"), optimizer.get_slot(weight, "mask"))+    masked_expected = tf.math.multiply(expected_saved_initialization, optimizer.get_slot(weight, "mask"))+    self.assertAllEqual(initialization_slot_k, masked_expected)++    mask_after_pruning = optimizer.get_slot(weight, "mask").read_value()+    self.assertAllEqual(np.count_nonzero(mask_after_pruning), 97)++  def testReloadWeightsatInitialization(self):+    weight1 = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")++    save_round = 0+    n_rounds = 20+    end_iter = 100+    frequency, prune_ratio_per_round = get_lth_sparsity(save_round, n_rounds, self.target_sparsity, end_iter)+    pruning_schedule = make_pruning_schedule(1 - prune_ratio_per_round, save_round, end_iter, frequency)+    +    p = pruner.LTHPruner(+        pruning_schedule=pruning_schedule,+        save_iteration=self.save_init,+        block_size=self.block_size,+        block_pooling_type=self.block_pooling_type)++    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)+    expected_saved_initialization = None++    def _train(weight):+      expected_saved_initialization = tf.ones_like(weight, dtype=weight.dtype.base_dtype) * -1+      p.create_slots(optimizer, weight)+      for i in tf.range(0, 2):+        p.preprocess_weights(optimizer, weight, self.grad(weight))+        if i == save_round:+          expected_saved_initialization = weight.read_value()+        weight.assign(tf.math.add(weight, sample_noise(i)))+        p.postprocess_weights(optimizer, weight, self.grad(weight))+        optimizer.iterations.assign_add(1) # save weights right before iteration+      return expected_saved_initialization++    expected_saved_initialization = _train(weight1)+    initialization_slot = optimizer.get_slot(weight1, "original_initialization")+    self.assertAllEqual(initialization_slot, expected_saved_initialization)++    mask_after_pruning = optimizer.get_slot(weight1, "mask").read_value()+    masked_weight_expected = tf.math.multiply(initialization_slot, mask_after_pruning)+    self.assertAllEqual(np.count_nonzero(masked_weight_expected), 97)+    self.assertAllEqual(np.count_nonzero(mask_after_pruning), 97)++    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")+    optimizer.iterations.assign(0)+    expected_saved_initialization = tf.function(_train)(weight)+    initialization_slot = optimizer.get_slot(weight, "original_initialization")+    self.assertAllEqual(initialization_slot, expected_saved_initialization)++    mask_after_pruning = optimizer.get_slot(weight, "mask").read_value()+    masked_weight_expected = tf.math.multiply(initialization_slot, mask_after_pruning)+    self.assertAllEqual(np.count_nonzero(masked_weight_expected), 97)+    self.assertAllEqual(np.count_nonzero(mask_after_pruning), 97)++  +  def testReloadWeightsIterationK(self):+    weight = tf.Variable(np.linspace(1.0, 100.0, 100), name="weights")++    save_round = 5+    n_rounds = 24+    end_iter = 100+    frequency, prune_ratio_per_round = get_lth_sparsity(save_round, n_rounds, self.target_sparsity, end_iter)+    pruning_schedule = make_pruning_schedule(1 - prune_ratio_per_round, save_round + 1, end_iter, frequency)+    +    p = pruner.LTHPruner(+        pruning_schedule=pruning_schedule,+        save_iteration=save_round,+        block_size=self.block_size,+        block_pooling_type=self.block_pooling_type)++    optimizer = self.dummy_optimizer+    optimizer.iterations.assign(0)++    def _train(weight):