Assign weight 1.0 to only examples from the given task.

def weights_multi_problem_all(labels, taskid=-1):
  """Assign weight 1.0 to only examples from the given task."""
  taskid = check_nonnegative(taskid)
  weights = to_float(tf.not_equal(labels, 0))
  past_taskid = tf.cumsum(to_float(tf.equal(labels, taskid)), axis=1)
  # Additionally zero out the task id location
  past_taskid *= to_float(tf.not_equal(labels, taskid))
  non_taskid = to_float(labels)
  example_mask = to_float(tf.not_equal(past_taskid * non_taskid, 0))
  example_mask = tf.reduce_sum(example_mask, axis=1)
  example_mask = to_float(
      tf.greater(example_mask, tf.zeros_like(example_mask)))

  return weights * tf.expand_dims(example_mask, axis=-1)