class_docstr: Implementation of magnitude-based weight pruning.
class_docstr_tok: ["Implementation","of","magnitude-based","weight","pruning","."]
docstr: Gather the assign objs for assigning weights<=weights*mask.
  
  The objs are ops for graph execution and tensors for eager
  execution.
  
  Returns:
    group of objs for weight assignment.
docstr_tok: ["Gather","the","assign","objs","for","assigning","weights","<","=weights","*","mask",".","The","objs","are","ops","for","graph","execution","and","tensors","for","eager","execution",".","Returns",":","group","of","objs","for","weight","assignment","."]
code: def _weight_assign_objs(self):
      """Gather the assign objs for assigning weights<=weights*mask.
  
      The objs are ops for graph execution and tensors for eager
      execution.
  
      Returns:
        group of objs for weight assignment.
      """
  
      def update_fn(distribution, values_and_vars):
          # TODO(yunluli): Need this ReduceOp because the weight is created by the
          # layer wrapped, so we don't have control of its aggregation policy. May
          # be able to optimize this when distribution strategy supports easier
          # update to mirrored variables in replica context.
          reduced_values = distribution.extended.batch_reduce_to(
              tf.distribute.ReduceOp.MEAN, values_and_vars
          )
          var_list = [v for _, v in values_and_vars]
          values_and_vars = zip(reduced_values, var_list)
  
          def update_var(variable, reduced_value):
              return tf_compat.assign(variable, reduced_value)
  
          update_objs = []
          for value, var in values_and_vars:
              update_objs.append(
                  distribution.extended.update(
                      var, update_var, args=(value,)
                  )
              )
  
          return tf.group(update_objs)
  
      assign_objs = []
  
      if tf.distribute.get_replica_context():
          values_and_vars = []
          for weight, mask, _ in self._pruning_vars:
              masked_weight = tf.dtypes.cast(
                  tf.math.multiply(weight, mask), dtype=weight.dtype
              )
              values_and_vars.append((masked_weight, weight))
          if values_and_vars:
              assign_objs.append(
                  tf.distribute.get_replica_context().merge_call(
                      update_fn, args=(values_and_vars,)
                  )
              )
      else:
          for weight, mask, _ in self._pruning_vars:
              masked_weight = tf.dtypes.cast(
                  tf.math.multiply(weight, mask), dtype=weight.dtype
              )
              assign_objs.append(
                  tf_compat.assign(weight, masked_weight)
              )
  
      return assign_objs
  
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 6 for calling the Pruning obj's underlying member method code with example usage: obj._weight_assign_objs() and returns: assign_objs, tf, tf_compat