class_docstr: Implementation of magnitude-based weight pruning.
class_docstr_tok: ["Implementation","of","magnitude-based","weight","pruning","."]
docstr: Returns an op to updates masks as per the pruning schedule.
docstr_tok: ["Returns","an","op","to","updates","masks","as","per","the","pruning","schedule","."]
code: def conditional_mask_update(self):
      """Returns an op to updates masks as per the pruning schedule."""
  
      def maybe_update_masks():
          if self._sparsity_m_by_n:
              # Update structured sparsity masks only at step 1
              return tf.math.equal(self._step_fn(), 1)
          else:
              return self._pruning_schedule(self._step_fn())[0]
  
      def no_update():
          return tf.no_op()
  
      def mask_update():
          """Updates mask without distribution strategy."""
  
          def update():
              assign_objs = []
  
              for weight, mask, threshold in self._pruning_vars:
                  (
                      new_threshold,
                      new_mask,
                  ) = self._maybe_update_block_mask(weight)
                  assign_objs.append(
                      tf_compat.assign(threshold, new_threshold)
                  )
                  assign_objs.append(tf_compat.assign(mask, new_mask))
  
              return tf.group(assign_objs)
  
          return tf.cond(maybe_update_masks(), update, no_update)
  
      def mask_update_distributed(distribution):
          """Updates mask with distribution strategy."""
  
          def update(var, value):
              return tf_compat.assign(var, value)
  
          def update_distributed():
              """Gather distributed update objs.
  
              The objs are ops for graph execution and tensors for eager
              execution.
              """
              assign_objs = []
  
              for weight, mask, threshold in self._pruning_vars:
                  (
                      new_threshold,
                      new_mask,
                  ) = self._maybe_update_block_mask(weight)
                  assign_objs.append(
                      distribution.extended.update(
                          mask, update, (new_mask,)
                      )
                  )
                  assign_objs.append(
                      distribution.extended.update(
                          threshold, update, (new_threshold,)
                      )
                  )
  
              return tf.group(assign_objs)
  
          return tf.cond(
              maybe_update_masks(), update_distributed, no_update
          )
  
      if tf.distribute.get_replica_context():
          return tf.distribute.get_replica_context().merge_call(
              mask_update_distributed
          )
      else:
          return mask_update()
  
file_path: model-optimization/tensorflow_model_optimization/python/core/sparsity/keras/pruning_impl.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import tensorflow
  from tensorflow_model_optimization.python.core.keras import compat
  from tensorflow_model_optimization.python.core.sparsity.keras import pruning_utils
label_desc: Use image node_id 8 for calling the Pruning obj's underlying member method code with example usage: obj.conditional_mask_update() and returns: tf, tf, tf, tf, mask_update, tf, self, tf, tf_compat, tf