//  [python]
def get_avuc_loss(probabilities, labels, avuc_stop_prob_gradients,
                  avuc_entropy_threshold):
  """Computes and returns the AvUC loss tensor.

  AvUC loss is defined in equation (3) in this paper:
  https://arxiv.org/pdf/2012.07923.pdf. AvUC-with-gradient-stopping is defined
  in the appendix of this paper: https://arxiv.org/pdf/2108.00106.pdf.

  Args:
    probabilities: tensor of predicted probabilities of
      (batch-size, num-classes) shape
    labels: tensor of labels in [0,num-classes) of (batch-size,) shape
    avuc_stop_prob_gradients: whether to use gradient-stopping
    avuc_entropy_threshold: entropy threshold (u_th in equation (3) cited above)

  Returns:
    A tensor of () shape containing a single value: the AvUC loss.
  """

  if avuc_stop_prob_gradients:
    confidences = tf.stop_gradient(tf.reduce_max(probabilities, axis=1))
  else:
    confidences = tf.reduce_max(probabilities, axis=1)

  accuracies = tf.cast(
      tf.math.equal(
          tf.argmax(probabilities, axis=1),
          tf.cast(labels, tf.int64),
      ),
      tf.float32,
  )

  dim = int(probabilities.shape[1])
  uniform_probabilities = tf.convert_to_tensor([1.0 / dim] * dim)
  log_safe_probabilities = ((1.0 - EPS) * probabilities +
                            EPS * uniform_probabilities)
  log_probabilities = tf.math.log(log_safe_probabilities)
  entropies = tf.math.negative(
      tf.reduce_sum(
          tf.multiply(log_safe_probabilities, log_probabilities), axis=1))

  accuracies_entropies_and_confidences = tf.stack(
      [accuracies, entropies, confidences], axis=-1)

  ac_clause = (lambda aec: aec[1] < avuc_entropy_threshold and aec[0] > 0.5)
  au_clause = (lambda aec: aec[1] >= avuc_entropy_threshold and aec[0] > 0.5)
  ic_clause = (lambda aec: aec[1] < avuc_entropy_threshold and aec[0] < 0.5)
  iu_clause = (lambda aec: aec[1] >= avuc_entropy_threshold and aec[0] < 0.5)

  # pylint: disable=g-long-lambda
  nac_diff = tf.reduce_sum(
      tf.map_fn(
          elems=accuracies_entropies_and_confidences,
          fn=lambda aec: tf.convert_to_tensor(aec[2]) *
          (tf.constant(1.0) - tf.math.tanh(tf.convert_to_tensor(aec[1])))
          if ac_clause(aec) else 0.0))
  nau_diff = tf.reduce_sum(
      tf.map_fn(
          elems=accuracies_entropies_and_confidences,
          fn=lambda aec: tf.convert_to_tensor(aec[2]) * tf.math.tanh(
              tf.convert_to_tensor(aec[1])) if au_clause(aec) else 0.0))
  nic_diff = tf.reduce_sum(
      tf.map_fn(
          elems=accuracies_entropies_and_confidences,
          fn=lambda aec: (tf.constant(1.0) - tf.convert_to_tensor(aec[2])) *
          (tf.constant(1.0) - tf.math.tanh(tf.convert_to_tensor(aec[1])))
          if ic_clause(aec) else 0.0))
  niu_diff = tf.reduce_sum(
      tf.map_fn(
          elems=accuracies_entropies_and_confidences,
          fn=lambda aec:
          (tf.constant(1.0) - tf.convert_to_tensor(aec[2])) * tf.math.tanh(
              tf.convert_to_tensor(aec[1])) if iu_clause(aec) else 0.0))
  # pylint: enable=g-long-lambda

  avuc_loss = tf.math.log(
      tf.constant(1.0) +
      ((nau_diff + nic_diff) /
       tf.math.maximum(nac_diff + niu_diff, tf.constant(EPS))))

  return avuc_loss