def class_logit(layer, label):
  """Like channel, but for softmax layers.

  Args:
    layer: A layer name string.
    label: Either a string (refering to a label in model.labels) or an int
      label position.

  Returns:
    Objective maximizing a logit.
  """
  def inner(T):
    if isinstance(label, int):
      class_n = label
    else:
      class_n = T("labels").index(label)
    logits = T(layer)
    logit = tf.reduce_sum(logits[:, class_n])
    return logit
  return inner