Weight-level magnitude pruning.

def weight_targeting(w, k):
  """Weight-level magnitude pruning."""
  k = tf.to_int32(k)
  w_shape = shape_list(w)
  size = tf.to_int32(tf.reduce_prod(w_shape[:-1]))
  w = tf.reshape(w, [size, w_shape[-1]])

  transpose_w = tf.transpose(w)
  thres = tf.contrib.framework.sort(tf.abs(transpose_w), axis=1)[:, k]
  mask = to_float(thres[None, :] >= tf.abs(w))

  return tf.reshape(mask, w_shape)