import tensorflow as tf

tf.app.flags.DEFINE_float('ws_iter_ratio_end', 0.5, 'WS: iteration ratio (at ending time)')
tf.app.flags.DEFINE_float('ws_mask_update_step', 500, 'WS: step size for updating the pruning mask')

def calc_prune_ratio(vars_list):
  """Calculate the overall pruning ratio for the given list of variables.
  Args:
  * vars_list: list of variables
  Returns:
  * prune_ratio: overall pruning ratio of the given list of variables
  """

  nb_params_nnz = tf.add_n([tf.count_nonzero(var) for var in vars_list])
  nb_params_all = tf.add_n([tf.size(var) for var in vars_list])
  prune_ratio = 1.0 - tf.cast(nb_params_nnz, tf.float32) / tf.cast(nb_params_all, tf.float32)

  return prune_ratio

class WeightSparseLearner(AbstractLearner):  # pylint: disable=too-many-instance-attributes
  """Weight sparsification learner."""

  def __init__(self, sm_writer, model_helper):
    """Constructor function.
    Args:
    * sm_writer: TensorFlow's summary writer
