parameters: (self, weights)
return_statement: return current_threshold, new_mask
docstring: 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
docstring_summary: Updates the mask for a given weight tensor.
docstring_tokens: Updates the mask for a given weight tensor .
function: 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.round(
                tf.dtypes.cast(tf.size(abs_weights), tf.float32) *
                (1 - sparsity)), tf.int32)
        # Sort the entire array
        values, _ = tf.math.top_k(
            tf.reshape(abs_weights, [-1]), k=tf.size(abs_weights))
        # Grab the (k-1)th value
  
        current_threshold = tf.gather(values, k - 1)
        new_mask = tf.dtypes.cast(
            tf.math.greater_equal(abs_weights, current_threshold), weights.dtype)
      return current_threshold, new_mask
function_tokens: def _update_mask ( self , weights ) : 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 . round ( tf . dtypes . cast ( tf . size ( abs_weights ) , tf . float32 ) * ( 1 - sparsity ) ) , tf . int32 ) # Sort the entire array values , _ = tf . math . top_k ( tf . reshape ( abs_weights , [ - 1 ] ) , k = tf . size ( abs_weights ) ) # Grab the (k-1)th value current_threshold = tf . gather ( values , k - 1 ) new_mask = tf . dtypes . cast ( tf . math . greater_equal ( abs_weights , current_threshold ) , weights . dtype ) return current_threshold , new_mask