class_docstr: Implementation of magnitude-based weight pruning.
class_docstr_tok: ["Implementation","of","magnitude-based","weight","pruning","."]
docstr: Updates the mask for a given weight tensor.
  
  This functions first estimates the threshold value such that
  a given fraction of weights have magnitude less than
  the threshold.
  
  Args:
    weights: The weight tensor that needs to be masked.
  
  Returns:
    new_threshold: The new value of the threshold based on weights, and
      sparsity at the current global_step
    new_mask: A numpy array of the same size and shape as weights containing
      0 or 1 to indicate which of the values in weights falls below
      the threshold
  
  Raises:
    ValueError: if sparsity is not defined
docstr_tok: ["Updates","the","mask","for","a","given","weight","tensor",".","This","functions","first","estimates","the","threshold","value","such","that","a","given","fraction","of","weights","have","magnitude","less","than","the","threshold",".","Args",":","weights",":","The","weight","tensor","that","needs","to","be","masked",".","Returns",":","new_threshold",":","The","new","value","of","the","threshold","based","on","weights",",","and","sparsity","at","the","current","global_step","new_mask",":","A","numpy","array","of","the","same","size","and","shape","as","weights","containing","0","or","1","to","indicate","which","of","the","values","in","weights","falls","below","the","threshold","Raises",":","ValueError",":","if","sparsity","is","not","defined"]
code: def _update_mask(self, weights):
      """Updates the mask for a given weight tensor.
  
      This functions first estimates the threshold value such that
      a given fraction of weights have magnitude less than
      the threshold.
  
      Args:
        weights: The weight tensor that needs to be masked.
  
      Returns:
        new_threshold: The new value of the threshold based on weights, and
          sparsity at the current global_step
        new_mask: A numpy array of the same size and shape as weights containing
          0 or 1 to indicate which of the values in weights falls below
          the threshold
  
      Raises:
        ValueError: if sparsity is not defined
      """
      sparsity = self._pruning_schedule(self._step_fn())[1]
      with tf.name_scope("pruning_ops"):
          abs_weights = tf.math.abs(weights)
          k = tf.dtypes.cast(
              tf.math.maximum(
                  tf.math.round(
                      tf.dtypes.cast(tf.size(abs_weights), tf.float32)
                      * (1 - sparsity)
                  ),
                  1,
              ),
              tf.int32,
          )
          values, indices = tf.math.top_k(
              tf.reshape(abs_weights, [-1]), k=tf.size(abs_weights)
          )
  
          # Grab the (k-1)th value as a threshold to build pruning mask.
          threshold_value = tf.gather(values, k - 1)
          threshold_pos = tf.gather(indices, k - 1)
  
          # Build mask for the weight element higher magnitude than threshold value.
          # A mask is added to make sure the threshold element be incorporated.
          # TODO(b/208967539): Update the logic to index oriented logic.
          new_mask = tf.math.logical_or(
              tf.math.greater_equal(abs_weights, threshold_value),
              tf.reshape(
                  tf.one_hot(
                      threshold_pos,
                      depth=tf.size(abs_weights),
                      on_value=True,
                      off_value=False,
                      dtype=tf.bool,
                  ),
                  abs_weights.shape,
              ),
          )
  
      # Updated mask is casted back to weight's data type in case of the type
      # mismatching due to keras mixed precision policy.
      return tf.dtypes.cast(
          threshold_value, weights.dtype
      ), tf.dtypes.cast(new_mask, weights.dtype)
  
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 3 for calling the Pruning obj's underlying member method code with example usage: obj._update_mask(weights) and returns: tf,  tf