We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Pruning: Training with Near 100% Target Sparsity Fails
**Describe the bug**
Pruning with high target sparsity (e.g. 0.99) causes a error.

**System information**

TensorFlow installed from (source or binary):

TensorFlow version: any

TensorFlow Model Optimization version: 0.2.1

Python version: any

**Describe the expected behavior**
Target sparsity of 0.99 should work. 

**Describe the current behavior**
Training errors out with something like:

InvalidArgumentError: indices = -1 is not in [0, 40)
	 [[{{node prune_low_magnitude_dense_1/cond/cond/pruning_ops/GatherV2}}]]

**Code to reproduce the issue**
testPruneWithHighSparsity_Fails in prune_integration_test.py

Can search for "model-optimization/issues/215" in codebase to find unit test also.



--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `tensorflow_model_optimization/python/core/sparsity/keras/pruning_impl.py`
Content:
```
1 # Copyright 2019 The TensorFlow Authors. All Rights Reserved.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 # ==============================================================================
15 """Helper functions to add support for magnitude-based model pruning."""
16 
17 from __future__ import absolute_import
18 from __future__ import division
19 from __future__ import print_function
20 
21 import tensorflow as tf
22 
23 from tensorflow.python.ops import summary_ops_v2
24 from tensorflow.python.summary import summary as summary_ops_v1
25 from tensorflow_model_optimization.python.core.keras import compat as tf_compat
26 from tensorflow_model_optimization.python.core.sparsity.keras import pruning_utils
27 
28 
29 class Pruning(object):
30   """Implementation of magnitude-based weight pruning."""
31 
32   def __init__(self, training_step_fn, pruning_vars, pruning_schedule,
33                block_size, block_pooling_type):
34     """The logic for magnitude-based pruning weight tensors.
35 
36     Args:
37       training_step_fn: A callable that returns the training step.
38       pruning_vars: A list of (weight, mask, threshold) tuples
39       pruning_schedule: A `PruningSchedule` object that controls pruning rate
40         throughout training.
41       block_size: The dimensions (height, weight) for the block sparse pattern
42         in rank-2 weight tensors.
43       block_pooling_type: (optional) The function to use to pool weights in the
44         block. Must be 'AVG' or 'MAX'.
45     """
46     self._pruning_vars = pruning_vars
47     self._pruning_schedule = pruning_schedule
48     self._block_size = list(block_size)
49     self._block_pooling_type = block_pooling_type
50     self._validate_block()
51 
52     # Training step
53     self._step_fn = training_step_fn
54 
55     self._validate_block()
56 
57   def _validate_block(self):
58     if self._block_size != [1, 1]:
59       for weight, _, _ in self._pruning_vars:
60         if weight.get_shape().ndims != 2:
61           raise ValueError('Block Sparsity can only be used for layers which '
62                            'have 2-dimensional weights.')
63 
64   def _update_mask(self, weights):
65     """Updates the mask for a given weight tensor.
66 
67     This functions first estimates the threshold value such that
68     a given fraction of weights have magnitude less than
69     the threshold.
70 
71     Args:
72       weights: The weight tensor that needs to be masked.
73 
74     Returns:
75       new_threshold: The new value of the threshold based on weights, and
76         sparsity at the current global_step
77       new_mask: A numpy array of the same size and shape as weights containing
78         0 or 1 to indicate which of the values in weights falls below
79         the threshold
80 
81     Raises:
82       ValueError: if sparsity is not defined
83     """
84     sparsity = self._pruning_schedule(self._step_fn())[1]
85     with tf.name_scope('pruning_ops'):
86       abs_weights = tf.math.abs(weights)
87       k = tf.dtypes.cast(
88           tf.math.round(
89               tf.dtypes.cast(tf.size(abs_weights), tf.float32) *
90               (1 - sparsity)), tf.int32)
91       # Sort the entire array
92       values, _ = tf.math.top_k(
93           tf.reshape(abs_weights, [-1]), k=tf.size(abs_weights))
94       # Grab the (k-1)th value
95 
96       current_threshold = tf.gather(values, k - 1)
97       new_mask = tf.dtypes.cast(
98           tf.math.greater_equal(abs_weights, current_threshold), weights.dtype)
99     return current_threshold, new_mask
100 
101   def _maybe_update_block_mask(self, weights):
102     """Performs block-granular masking of the weights.
103 
104     Block pruning occurs only if the block_height or block_width is > 1 and
105     if the weight tensor, when squeezed, has ndims = 2. Otherwise, elementwise
106     pruning occurs.
107     Args:
108       weights: The weight tensor that needs to be masked.
109 
110     Returns:
111       new_threshold: The new value of the threshold based on weights, and
112         sparsity at the current global_step
113       new_mask: A numpy array of the same size and shape as weights containing
114         0 or 1 to indicate which of the values in weights falls below
115         the threshold
116 
117     Raises:
118       ValueError: if block pooling function is not AVG or MAX
119     """
120     if self._block_size == [1, 1]:
121       return self._update_mask(weights)
122 
123     # TODO(pulkitb): Check if squeeze operations should now be removed since
124     # we are only accepting 2-D weights.
125 
126     squeezed_weights = tf.squeeze(weights)
127     abs_weights = tf.math.abs(squeezed_weights)
128     pooled_weights = pruning_utils.factorized_pool(
129         abs_weights,
130         window_shape=self._block_size,
131         pooling_type=self._block_pooling_type,
132         strides=self._block_size,
133         padding='SAME')
134 
135     if pooled_weights.get_shape().ndims != 2:
136       pooled_weights = tf.squeeze(pooled_weights)
137 
138     new_threshold, new_mask = self._update_mask(pooled_weights)
139 
140     updated_mask = pruning_utils.expand_tensor(new_mask, self._block_size)
141     sliced_mask = tf.slice(
142         updated_mask, [0, 0],
143         [squeezed_weights.get_shape()[0],
144          squeezed_weights.get_shape()[1]])
145     return new_threshold, tf.reshape(sliced_mask, tf.shape(weights))
146 
147   def _weight_assign_objs(self):
148     """Gather the assign objs for assigning weights<=weights*mask.
149 
150     The objs are ops for graph execution and tensors for eager
151     execution.
152 
153     Returns:
154       group of objs for weight assignment.
155     """
156 
157     def update_fn(distribution, values_and_vars):
158       # TODO(yunluli): Need this ReduceOp because the weight is created by the
159       # layer wrapped, so we don't have control of its aggregation policy. May
160       # be able to optimize this when distribution strategy supports easier
161       # update to mirrored variables in replica context.
162       reduced_values = distribution.extended.batch_reduce_to(
163           tf.distribute.ReduceOp.MEAN, values_and_vars)
164       var_list = [v for _, v in values_and_vars]
165       values_and_vars = zip(reduced_values, var_list)
166 
167       def update_var(variable, reduced_value):
168         return tf_compat.assign(variable, reduced_value)
169 
170       update_objs = []
171       for value, var in values_and_vars:
172         update_objs.append(
173             distribution.extended.update(var, update_var, args=(value,)))
174 
175       return tf.group(update_objs)
176 
177     assign_objs = []
178 
179     if tf.distribute.get_replica_context():
180       values_and_vars = []
181       for weight, mask, _ in self._pruning_vars:
182         masked_weight = tf.math.multiply(weight, mask)
183         values_and_vars.append((masked_weight, weight))
184       if values_and_vars:
185         assign_objs.append(tf.distribute.get_replica_context().merge_call(
186             update_fn, args=(values_and_vars,)))
187     else:
188       for weight, mask, _ in self._pruning_vars:
189         masked_weight = tf.math.multiply(weight, mask)
190         assign_objs.append(tf_compat.assign(weight, masked_weight))
191 
192     return assign_objs
193 
194   def weight_mask_op(self):
195     return tf.group(self._weight_assign_objs())
196 
197   def conditional_mask_update(self):
198     """Returns an op to updates masks as per the pruning schedule."""
199 
200     def maybe_update_masks():
201       return self._pruning_schedule(self._step_fn())[0]
202 
203     def no_update():
204       return tf.no_op()
205 
206     def mask_update():
207       """Updates mask without distribution strategy."""
208 
209       def update():
210         assign_objs = []
211 
212         for weight, mask, threshold in self._pruning_vars:
213           new_threshold, new_mask = self._maybe_update_block_mask(weight)
214           assign_objs.append(tf_compat.assign(threshold, new_threshold))
215           assign_objs.append(tf_compat.assign(mask, new_mask))
216 
217         return tf.group(assign_objs)
218 
219       return tf.cond(maybe_update_masks(), update, no_update)
220 
221     def mask_update_distributed(distribution):
222       """Updates mask with distribution strategy."""
223 
224       def update(var, value):
225         return tf_compat.assign(var, value)
226 
227       def update_distributed():
228         """Gather distributed update objs.
229 
230         The objs are ops for graph execution and tensors for eager
231         execution.
232         """
233         assign_objs = []
234 
235         for weight, mask, threshold in self._pruning_vars:
236           new_threshold, new_mask = self._maybe_update_block_mask(weight)
237           assign_objs.append(
238               distribution.extended.update(mask, update, (new_mask,)))
239           assign_objs.append(
240               distribution.extended.update(threshold, update, (new_threshold,)))
241 
242         return tf.group(assign_objs)
243 
244       return tf.cond(maybe_update_masks(), update_distributed, no_update)
245 
246     if tf.distribute.get_replica_context():
247       return tf.distribute.get_replica_context().merge_call(
248           mask_update_distributed)
249     else:
250       return mask_update()
251 
252   def add_pruning_summaries(self):
253     """Adds summaries of weight sparsities and thresholds."""
254     # b/(139939526): update to use public API.
255     summary = summary_ops_v1
256     if tf.executing_eagerly():
257       summary = summary_ops_v2
258     summary.scalar('sparsity', self._pruning_schedule(self._step_fn())[1])
259     for _, mask, threshold in self._pruning_vars:
260       summary.scalar(mask.name + '/sparsity', 1.0 - tf.math.reduce_mean(mask))
261       summary.scalar(threshold.name + '/threshold', threshold)
262
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/tensorflow_model_optimization/python/core/sparsity/keras/pruning_impl.py b/tensorflow_model_optimization/python/core/sparsity/keras/pruning_impl.py
--- a/tensorflow_model_optimization/python/core/sparsity/keras/pruning_impl.py
+++ b/tensorflow_model_optimization/python/core/sparsity/keras/pruning_impl.py
@@ -85,9 +85,12 @@
     with tf.name_scope('pruning_ops'):
       abs_weights = tf.math.abs(weights)
       k = tf.dtypes.cast(
-          tf.math.round(
-              tf.dtypes.cast(tf.size(abs_weights), tf.float32) *
-              (1 - sparsity)), tf.int32)
+          tf.math.maximum(
+              tf.math.round(
+                  tf.dtypes.cast(tf.size(abs_weights), tf.float32) *
+                  (1 - sparsity)),
+              1),
+          tf.int32)
       # Sort the entire array
       values, _ = tf.math.top_k(
           tf.reshape(abs_weights, [-1]), k=tf.size(abs_weights))

{"golden_diff": "diff --git a/tensorflow_model_optimization/python/core/sparsity/keras/pruning_impl.py b/tensorflow_model_optimization/python/core/sparsity/keras/pruning_impl.py\n--- a/tensorflow_model_optimization/python/core/sparsity/keras/pruning_impl.py\n+++ b/tensorflow_model_optimization/python/core/sparsity/keras/pruning_impl.py\n@@ -85,9 +85,12 @@\n     with tf.name_scope('pruning_ops'):\n       abs_weights = tf.math.abs(weights)\n       k = tf.dtypes.cast(\n-          tf.math.round(\n-              tf.dtypes.cast(tf.size(abs_weights), tf.float32) *\n-              (1 - sparsity)), tf.int32)\n+          tf.math.maximum(\n+              tf.math.round(\n+                  tf.dtypes.cast(tf.size(abs_weights), tf.float32) *\n+                  (1 - sparsity)),\n+              1),\n+          tf.int32)\n       # Sort the entire array\n       values, _ = tf.math.top_k(\n           tf.reshape(abs_weights, [-1]), k=tf.size(abs_weights))\n", "issue": "Pruning: Training with Near 100% Target Sparsity Fails\n**Describe the bug**\r\nPruning with high target sparsity (e.g. 0.99) causes a error.\r\n\r\n**System information**\r\n\r\nTensorFlow installed from (source or binary):\r\n\r\nTensorFlow version: any\r\n\r\nTensorFlow Model Optimization version: 0.2.1\r\n\r\nPython version: any\r\n\r\n**Describe the expected behavior**\r\nTarget sparsity of 0.99 should work. \r\n\r\n**Describe the current behavior**\r\nTraining errors out with something like:\r\n\r\nInvalidArgumentError: indices = -1 is not in [0, 40)\r\n\t [[{{node prune_low_magnitude_dense_1/cond/cond/pruning_ops/GatherV2}}]]\r\n\r\n**Code to reproduce the issue**\r\ntestPruneWithHighSparsity_Fails in prune_integration_test.py\r\n\r\nCan search for \"model-optimization/issues/215\" in codebase to find unit test also.\r\n\r\n\n", "before_files": [{"content": "# Copyright 2019 The TensorFlow Authors. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ==============================================================================\n\"\"\"Helper functions to add support for magnitude-based model pruning.\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport tensorflow as tf\n\nfrom tensorflow.python.ops import summary_ops_v2\nfrom tensorflow.python.summary import summary as summary_ops_v1\nfrom tensorflow_model_optimization.python.core.keras import compat as tf_compat\nfrom tensorflow_model_optimization.python.core.sparsity.keras import pruning_utils\n\n\nclass Pruning(object):\n  \"\"\"Implementation of magnitude-based weight pruning.\"\"\"\n\n  def __init__(self, training_step_fn, pruning_vars, pruning_schedule,\n               block_size, block_pooling_type):\n    \"\"\"The logic for magnitude-based pruning weight tensors.\n\n    Args:\n      training_step_fn: A callable that returns the training step.\n      pruning_vars: A list of (weight, mask, threshold) tuples\n      pruning_schedule: A `PruningSchedule` object that controls pruning rate\n        throughout training.\n      block_size: The dimensions (height, weight) for the block sparse pattern\n        in rank-2 weight tensors.\n      block_pooling_type: (optional) The function to use to pool weights in the\n        block. Must be 'AVG' or 'MAX'.\n    \"\"\"\n    self._pruning_vars = pruning_vars\n    self._pruning_schedule = pruning_schedule\n    self._block_size = list(block_size)\n    self._block_pooling_type = block_pooling_type\n    self._validate_block()\n\n    # Training step\n    self._step_fn = training_step_fn\n\n    self._validate_block()\n\n  def _validate_block(self):\n    if self._block_size != [1, 1]:\n      for weight, _, _ in self._pruning_vars:\n        if weight.get_shape().ndims != 2:\n          raise ValueError('Block Sparsity can only be used for layers which '\n                           'have 2-dimensional weights.')\n\n  def _update_mask(self, weights):\n    \"\"\"Updates the mask for a given weight tensor.\n\n    This functions first estimates the threshold value such that\n    a given fraction of weights have magnitude less than\n    the threshold.\n\n    Args:\n      weights: The weight tensor that needs to be masked.\n\n    Returns:\n      new_threshold: The new value of the threshold based on weights, and\n        sparsity at the current global_step\n      new_mask: A numpy array of the same size and shape as weights containing\n        0 or 1 to indicate which of the values in weights falls below\n        the threshold\n\n    Raises:\n      ValueError: if sparsity is not defined\n    \"\"\"\n    sparsity = self._pruning_schedule(self._step_fn())[1]\n    with tf.name_scope('pruning_ops'):\n      abs_weights = tf.math.abs(weights)\n      k = tf.dtypes.cast(\n          tf.math.round(\n              tf.dtypes.cast(tf.size(abs_weights), tf.float32) *\n              (1 - sparsity)), tf.int32)\n      # Sort the entire array\n      values, _ = tf.math.top_k(\n          tf.reshape(abs_weights, [-1]), k=tf.size(abs_weights))\n      # Grab the (k-1)th value\n\n      current_threshold = tf.gather(values, k - 1)\n      new_mask = tf.dtypes.cast(\n          tf.math.greater_equal(abs_weights, current_threshold), weights.dtype)\n    return current_threshold, new_mask\n\n  def _maybe_update_block_mask(self, weights):\n    \"\"\"Performs block-granular masking of the weights.\n\n    Block pruning occurs only if the block_height or block_width is > 1 and\n    if the weight tensor, when squeezed, has ndims = 2. Otherwise, elementwise\n    pruning occurs.\n    Args:\n      weights: The weight tensor that needs to be masked.\n\n    Returns:\n      new_threshold: The new value of the threshold based on weights, and\n        sparsity at the current global_step\n      new_mask: A numpy array of the same size and shape as weights containing\n        0 or 1 to indicate which of the values in weights falls below\n        the threshold\n\n    Raises:\n      ValueError: if block pooling function is not AVG or MAX\n    \"\"\"\n    if self._block_size == [1, 1]:\n      return self._update_mask(weights)\n\n    # TODO(pulkitb): Check if squeeze operations should now be removed since\n    # we are only accepting 2-D weights.\n\n    squeezed_weights = tf.squeeze(weights)\n    abs_weights = tf.math.abs(squeezed_weights)\n    pooled_weights = pruning_utils.factorized_pool(\n        abs_weights,\n        window_shape=self._block_size,\n        pooling_type=self._block_pooling_type,\n        strides=self._block_size,\n        padding='SAME')\n\n    if pooled_weights.get_shape().ndims != 2:\n      pooled_weights = tf.squeeze(pooled_weights)\n\n    new_threshold, new_mask = self._update_mask(pooled_weights)\n\n    updated_mask = pruning_utils.expand_tensor(new_mask, self._block_size)\n    sliced_mask = tf.slice(\n        updated_mask, [0, 0],\n        [squeezed_weights.get_shape()[0],\n         squeezed_weights.get_shape()[1]])\n    return new_threshold, tf.reshape(sliced_mask, tf.shape(weights))\n\n  def _weight_assign_objs(self):\n    \"\"\"Gather the assign objs for assigning weights<=weights*mask.\n\n    The objs are ops for graph execution and tensors for eager\n    execution.\n\n    Returns:\n      group of objs for weight assignment.\n    \"\"\"\n\n    def update_fn(distribution, values_and_vars):\n      # TODO(yunluli): Need this ReduceOp because the weight is created by the\n      # layer wrapped, so we don't have control of its aggregation policy. May\n      # be able to optimize this when distribution strategy supports easier\n      # update to mirrored variables in replica context.\n      reduced_values = distribution.extended.batch_reduce_to(\n          tf.distribute.ReduceOp.MEAN, values_and_vars)\n      var_list = [v for _, v in values_and_vars]\n      values_and_vars = zip(reduced_values, var_list)\n\n      def update_var(variable, reduced_value):\n        return tf_compat.assign(variable, reduced_value)\n\n      update_objs = []\n      for value, var in values_and_vars:\n        update_objs.append(\n            distribution.extended.update(var, update_var, args=(value,)))\n\n      return tf.group(update_objs)\n\n    assign_objs = []\n\n    if tf.distribute.get_replica_context():\n      values_and_vars = []\n      for weight, mask, _ in self._pruning_vars:\n        masked_weight = tf.math.multiply(weight, mask)\n        values_and_vars.append((masked_weight, weight))\n      if values_and_vars:\n        assign_objs.append(tf.distribute.get_replica_context().merge_call(\n            update_fn, args=(values_and_vars,)))\n    else:\n      for weight, mask, _ in self._pruning_vars:\n        masked_weight = tf.math.multiply(weight, mask)\n        assign_objs.append(tf_compat.assign(weight, masked_weight))\n\n    return assign_objs\n\n  def weight_mask_op(self):\n    return tf.group(self._weight_assign_objs())\n\n  def conditional_mask_update(self):\n    \"\"\"Returns an op to updates masks as per the pruning schedule.\"\"\"\n\n    def maybe_update_masks():\n      return self._pruning_schedule(self._step_fn())[0]\n\n    def no_update():\n      return tf.no_op()\n\n    def mask_update():\n      \"\"\"Updates mask without distribution strategy.\"\"\"\n\n      def update():\n        assign_objs = []\n\n        for weight, mask, threshold in self._pruning_vars:\n          new_threshold, new_mask = self._maybe_update_block_mask(weight)\n          assign_objs.append(tf_compat.assign(threshold, new_threshold))\n          assign_objs.append(tf_compat.assign(mask, new_mask))\n\n        return tf.group(assign_objs)\n\n      return tf.cond(maybe_update_masks(), update, no_update)\n\n    def mask_update_distributed(distribution):\n      \"\"\"Updates mask with distribution strategy.\"\"\"\n\n      def update(var, value):\n        return tf_compat.assign(var, value)\n\n      def update_distributed():\n        \"\"\"Gather distributed update objs.\n\n        The objs are ops for graph execution and tensors for eager\n        execution.\n        \"\"\"\n        assign_objs = []\n\n        for weight, mask, threshold in self._pruning_vars:\n          new_threshold, new_mask = self._maybe_update_block_mask(weight)\n          assign_objs.append(\n              distribution.extended.update(mask, update, (new_mask,)))\n          assign_objs.append(\n              distribution.extended.update(threshold, update, (new_threshold,)))\n\n        return tf.group(assign_objs)\n\n      return tf.cond(maybe_update_masks(), update_distributed, no_update)\n\n    if tf.distribute.get_replica_context():\n      return tf.distribute.get_replica_context().merge_call(\n          mask_update_distributed)\n    else:\n      return mask_update()\n\n  def add_pruning_summaries(self):\n    \"\"\"Adds summaries of weight sparsities and thresholds.\"\"\"\n    # b/(139939526): update to use public API.\n    summary = summary_ops_v1\n    if tf.executing_eagerly():\n      summary = summary_ops_v2\n    summary.scalar('sparsity', self._pruning_schedule(self._step_fn())[1])\n    for _, mask, threshold in self._pruning_vars:\n      summary.scalar(mask.name + '/sparsity', 1.0 - tf.math.reduce_mean(mask))\n      summary.scalar(threshold.name + '/threshold', threshold)\n", "path": "tensorflow_model_optimization/python/core/sparsity/keras/pruning_impl.py"}], "after_files": [{"content": "# Copyright 2019 The TensorFlow Authors. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ==============================================================================\n\"\"\"Helper functions to add support for magnitude-based model pruning.\"\"\"\n\nfrom __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport tensorflow as tf\n\nfrom tensorflow.python.ops import summary_ops_v2\nfrom tensorflow.python.summary import summary as summary_ops_v1\nfrom tensorflow_model_optimization.python.core.keras import compat as tf_compat\nfrom tensorflow_model_optimization.python.core.sparsity.keras import pruning_utils\n\n\nclass Pruning(object):\n  \"\"\"Implementation of magnitude-based weight pruning.\"\"\"\n\n  def __init__(self, training_step_fn, pruning_vars, pruning_schedule,\n               block_size, block_pooling_type):\n    \"\"\"The logic for magnitude-based pruning weight tensors.\n\n    Args:\n      training_step_fn: A callable that returns the training step.\n      pruning_vars: A list of (weight, mask, threshold) tuples\n      pruning_schedule: A `PruningSchedule` object that controls pruning rate\n        throughout training.\n      block_size: The dimensions (height, weight) for the block sparse pattern\n        in rank-2 weight tensors.\n      block_pooling_type: (optional) The function to use to pool weights in the\n        block. Must be 'AVG' or 'MAX'.\n    \"\"\"\n    self._pruning_vars = pruning_vars\n    self._pruning_schedule = pruning_schedule\n    self._block_size = list(block_size)\n    self._block_pooling_type = block_pooling_type\n    self._validate_block()\n\n    # Training step\n    self._step_fn = training_step_fn\n\n    self._validate_block()\n\n  def _validate_block(self):\n    if self._block_size != [1, 1]:\n      for weight, _, _ in self._pruning_vars:\n        if weight.get_shape().ndims != 2:\n          raise ValueError('Block Sparsity can only be used for layers which '\n                           'have 2-dimensional weights.')\n\n  def _update_mask(self, weights):\n    \"\"\"Updates the mask for a given weight tensor.\n\n    This functions first estimates the threshold value such that\n    a given fraction of weights have magnitude less than\n    the threshold.\n\n    Args:\n      weights: The weight tensor that needs to be masked.\n\n    Returns:\n      new_threshold: The new value of the threshold based on weights, and\n        sparsity at the current global_step\n      new_mask: A numpy array of the same size and shape as weights containing\n        0 or 1 to indicate which of the values in weights falls below\n        the threshold\n\n    Raises:\n      ValueError: if sparsity is not defined\n    \"\"\"\n    sparsity = self._pruning_schedule(self._step_fn())[1]\n    with tf.name_scope('pruning_ops'):\n      abs_weights = tf.math.abs(weights)\n      k = tf.dtypes.cast(\n          tf.math.maximum(\n              tf.math.round(\n                  tf.dtypes.cast(tf.size(abs_weights), tf.float32) *\n                  (1 - sparsity)),\n              1),\n          tf.int32)\n      # Sort the entire array\n      values, _ = tf.math.top_k(\n          tf.reshape(abs_weights, [-1]), k=tf.size(abs_weights))\n      # Grab the (k-1)th value\n\n      current_threshold = tf.gather(values, k - 1)\n      new_mask = tf.dtypes.cast(\n          tf.math.greater_equal(abs_weights, current_threshold), weights.dtype)\n    return current_threshold, new_mask\n\n  def _maybe_update_block_mask(self, weights):\n    \"\"\"Performs block-granular masking of the weights.\n\n    Block pruning occurs only if the block_height or block_width is > 1 and\n    if the weight tensor, when squeezed, has ndims = 2. Otherwise, elementwise\n    pruning occurs.\n    Args:\n      weights: The weight tensor that needs to be masked.\n\n    Returns:\n      new_threshold: The new value of the threshold based on weights, and\n        sparsity at the current global_step\n      new_mask: A numpy array of the same size and shape as weights containing\n        0 or 1 to indicate which of the values in weights falls below\n        the threshold\n\n    Raises:\n      ValueError: if block pooling function is not AVG or MAX\n    \"\"\"\n    if self._block_size == [1, 1]:\n      return self._update_mask(weights)\n\n    # TODO(pulkitb): Check if squeeze operations should now be removed since\n    # we are only accepting 2-D weights.\n\n    squeezed_weights = tf.squeeze(weights)\n    abs_weights = tf.math.abs(squeezed_weights)\n    pooled_weights = pruning_utils.factorized_pool(\n        abs_weights,\n        window_shape=self._block_size,\n        pooling_type=self._block_pooling_type,\n        strides=self._block_size,\n        padding='SAME')\n\n    if pooled_weights.get_shape().ndims != 2:\n      pooled_weights = tf.squeeze(pooled_weights)\n\n    new_threshold, new_mask = self._update_mask(pooled_weights)\n\n    updated_mask = pruning_utils.expand_tensor(new_mask, self._block_size)\n    sliced_mask = tf.slice(\n        updated_mask, [0, 0],\n        [squeezed_weights.get_shape()[0],\n         squeezed_weights.get_shape()[1]])\n    return new_threshold, tf.reshape(sliced_mask, tf.shape(weights))\n\n  def _weight_assign_objs(self):\n    \"\"\"Gather the assign objs for assigning weights<=weights*mask.\n\n    The objs are ops for graph execution and tensors for eager\n    execution.\n\n    Returns:\n      group of objs for weight assignment.\n    \"\"\"\n\n    def update_fn(distribution, values_and_vars):\n      # TODO(yunluli): Need this ReduceOp because the weight is created by the\n      # layer wrapped, so we don't have control of its aggregation policy. May\n      # be able to optimize this when distribution strategy supports easier\n      # update to mirrored variables in replica context.\n      reduced_values = distribution.extended.batch_reduce_to(\n          tf.distribute.ReduceOp.MEAN, values_and_vars)\n      var_list = [v for _, v in values_and_vars]\n      values_and_vars = zip(reduced_values, var_list)\n\n      def update_var(variable, reduced_value):\n        return tf_compat.assign(variable, reduced_value)\n\n      update_objs = []\n      for value, var in values_and_vars:\n        update_objs.append(\n            distribution.extended.update(var, update_var, args=(value,)))\n\n      return tf.group(update_objs)\n\n    assign_objs = []\n\n    if tf.distribute.get_replica_context():\n      values_and_vars = []\n      for weight, mask, _ in self._pruning_vars:\n        masked_weight = tf.math.multiply(weight, mask)\n        values_and_vars.append((masked_weight, weight))\n      if values_and_vars:\n        assign_objs.append(tf.distribute.get_replica_context().merge_call(\n            update_fn, args=(values_and_vars,)))\n    else:\n      for weight, mask, _ in self._pruning_vars:\n        masked_weight = tf.math.multiply(weight, mask)\n        assign_objs.append(tf_compat.assign(weight, masked_weight))\n\n    return assign_objs\n\n  def weight_mask_op(self):\n    return tf.group(self._weight_assign_objs())\n\n  def conditional_mask_update(self):\n    \"\"\"Returns an op to updates masks as per the pruning schedule.\"\"\"\n\n    def maybe_update_masks():\n      return self._pruning_schedule(self._step_fn())[0]\n\n    def no_update():\n      return tf.no_op()\n\n    def mask_update():\n      \"\"\"Updates mask without distribution strategy.\"\"\"\n\n      def update():\n        assign_objs = []\n\n        for weight, mask, threshold in self._pruning_vars:\n          new_threshold, new_mask = self._maybe_update_block_mask(weight)\n          assign_objs.append(tf_compat.assign(threshold, new_threshold))\n          assign_objs.append(tf_compat.assign(mask, new_mask))\n\n        return tf.group(assign_objs)\n\n      return tf.cond(maybe_update_masks(), update, no_update)\n\n    def mask_update_distributed(distribution):\n      \"\"\"Updates mask with distribution strategy.\"\"\"\n\n      def update(var, value):\n        return tf_compat.assign(var, value)\n\n      def update_distributed():\n        \"\"\"Gather distributed update objs.\n\n        The objs are ops for graph execution and tensors for eager\n        execution.\n        \"\"\"\n        assign_objs = []\n\n        for weight, mask, threshold in self._pruning_vars:\n          new_threshold, new_mask = self._maybe_update_block_mask(weight)\n          assign_objs.append(\n              distribution.extended.update(mask, update, (new_mask,)))\n          assign_objs.append(\n              distribution.extended.update(threshold, update, (new_threshold,)))\n\n        return tf.group(assign_objs)\n\n      return tf.cond(maybe_update_masks(), update_distributed, no_update)\n\n    if tf.distribute.get_replica_context():\n      return tf.distribute.get_replica_context().merge_call(\n          mask_update_distributed)\n    else:\n      return mask_update()\n\n  def add_pruning_summaries(self):\n    \"\"\"Adds summaries of weight sparsities and thresholds.\"\"\"\n    # b/(139939526): update to use public API.\n    summary = summary_ops_v1\n    if tf.executing_eagerly():\n      summary = summary_ops_v2\n    summary.scalar('sparsity', self._pruning_schedule(self._step_fn())[1])\n    for _, mask, threshold in self._pruning_vars:\n      summary.scalar(mask.name + '/sparsity', 1.0 - tf.math.reduce_mean(mask))\n      summary.scalar(threshold.name + '/threshold', threshold)\n", "path": "tensorflow_model_optimization/python/core/sparsity/keras/pruning_impl.py"}]}