code: def quantile_auc(q0, n0, q1, n1, curve="ROC", name=None):
      """Calculate ranking stats AUROC and AUPRC.
  
      Computes AUROC and AUPRC from quantiles (one for positive trials and one for
      negative trials).
  
      We use `pi(x)` to denote a score. We assume that if `pi(x) > k` for some
      threshold `k` then the event is predicted to be "1", and otherwise it is
      predicted to be a "0".  Its actual label is `y`, which may or may not be the
      same.
  
      Area Under Curve: Receiver Operator Characteristic (AUROC) is defined as:
  
      / 1
      |  TruePositiveRate(k) d FalsePositiveRate(k)
      / 0
  
      where,
  
      ```none
      TruePositiveRate(k) = P(pi(x) > k | y = 1)
      FalsePositiveRate(k) = P(pi(x) > k | y = 0)
      ```
  
      Area Under Curve: Precision-Recall (AUPRC) is defined as:
  
      / 1
      |  Precision(k) d Recall(k)
      / 0
  
      where,
  
      ```none
        Precision(k) = P(y = 1 | pi(x) > k)
        Recall(k) = TruePositiveRate(k) = P(pi(x) > k | y = 1)
      ```
  
      Notice that AUROC and AUPRC exchange the role of Recall in the
      integration, i.e.,
  
                  Integrand    Measure
                +------------+-----------+
        AUROC   |  Recall    |  FPR      |
                +------------+-----------+
        AUPRC   |  Precision |  Recall   |
                +------------+-----------+
  
      To learn more about the relationship between AUROC and AUPRC see [1].
  
      Args:
        q0: `N-D` `Tensor` of `float`, Quantiles of predicted probabilities given a
          negative trial. The first `N-1` dimensions are batch dimensions, and the
          AUC is calculated over the final dimension.
        n0: `float` or `(N-1)-D Tensor`, Number of negative trials. If `Tensor`,
          dimensions must match the first `N-1` dimensions of `q0`.
        q1: `N-D` `Tensor` of `float`, Quantiles of predicted probabilities given a
          positive trial. The first `N-1` dimensions are batch dimensions, which
          must match those of `q0`.
        n1: `float` or `(N-1)-D Tensor`, Number of positive trials. If `Tensor`,
          dimensions must match the first `N-1` dimensions of `q0`.
        curve: `str`, Specifies the name of the curve to be computed. Must be 'ROC'
          [default] or 'PR' for the Precision-Recall-curve.
        name: `str`, An optional name_scope name.
  
      Returns:
        auc: `Tensor` of `float`, area under the ROC or PR curve.
  
      #### Examples
  
      ```python
        n = 1000
        m = 500
        predictions_given_positive = np.random.rand(n)
        predictions_given_negative = np.random.rand(m)
        q1 = tfp.stats.quantiles(predictions_given_positive, num_quantiles=50)
        q0 = tfp.stats.quantiles(predictions_given_negative, num_quantiles=50)
        auroc = tfp.stats.quantile_auc(q0, m, q1, n, curve='ROC')
  
      ```
  
      ### Mathematical Details
  
      The algorithm proceeds by partitioning the combined quantile data into a
      series of intervals `[a, b)`, approximating the probability mass of
      predictions conditioned on a positive trial (`d1`) and probability mass of
      predictions conditioned on a negative trial (`d0`) in each interval `[a, b)`,
      and accumulating the incremental AUROC/AUPRC as functions of `d0` and `d1`.
  
      We assume that pi(x) is uniform within a given bucket of each quantile. Thus
      it will also be uniform within an interval [a, b) as long as the interval does
      not cross the quantile's bucket boundaries.
  
      A consequence of this assumption is that the cdf is piecewise linear.
      That is,
  
        P( pi(x) > k | y = 0 ), and,
        P( pi(x) > k | y = 1 ),
  
      are linear in `k`.
  
      Standard AUROC is fairly easier to calculate. Under the conditional uniformity
      assumptions we have a piece's contribution, [a, b), as:
  
       / b
       |
       |  P(y = 1 | pi > k) d P(pi > k | y = 0)
       |
       / a
  
          / b
          |
      = - |  P(pi > k | y = 1) P(pi = k | y = 0) d k
          |
          / a
  
                             / b
         -1 / (len(q0) - 1)  |
      = -------------------- |  P(pi > k | y = 1) d k
         q0[j + 1] - q0[j]   |
                             / a
  
                                    / 1
         1 / (len(q0) - 1)          |
      = ------------------- (b - a) |  (p1 + u d1) d u
         q0[j + 1] - q0[j]          |
                                    / 0
  
         1 / (len(q0) - 1)
      = ------------------- (b - a) (p1 + d1 / 2)
         q0[j + 1] - q0[j]
  
  
      AUPRC is a bit harder to calculate since the integrand,
      `P(y > 0 | pi(x) > k)`, is conditional on `k` rather than a probability over
      `k`.
  
      We proceed by formulating Precision in terms of the quantiles we have
      available to us.
  
      Precision(k) = P(y = 1 | pi(x) > k)
  
                         P(pi(x) > delta | y = 1 ) P(y = 1)
      = -----------------------------------------------------------------------
        P(pi(x) > delta | y = 1 ) P(y = 1) + P(pi(x) > delta | y = 0 ) P(y = 0)
  
  
      Since the cdf's are piecewise linear, we calculate this piece's contribution
      to AUPRC by the integral:
  
       / b
       |
       |  P(y = 1 | pi(x) > delta) d P(pi > delta | y = 1)
       |
       / a
  
         1 / (len(q1) - 1)
      = ------------------- (b - a) *
         q1[i + 1] - q1[i]
  
            / 1
            |             n1 * (u d1 + p1)
          * |  -------------------------------------- du
            |   n1 * (u d1 + p1) +  n0 * (u d0 + p0)
            / 0
  
                                    / 1
         1 / (len(q1) - 1)          |            n1 * (u d1 + p1)
      = ------------------- (b - a) |  ------------------------------------ du
         q1[i + 1] - q1[i]          |  n1 * (u d1 + p1) +  n0 * (u d0 + p0)
                                    / 0
  
      where the equality is a consequence of the piecewise uniformity assumption.
  
      The solution to the integral is given by Mathematica:
  
      ```
      Integrate[n1 (u d1 + p1) / (n1 (u d1 + p1) + n0 (u d0 + p0)), {u, 0, 1},
                Assumptions -> {p1 > 0, d1 > 0, p0 > 0, d0 > 0, n1 > 0, n0 > 0}]
      ```
  
      This integral can be solved by hand by noticing that:
  
        f(x) / (f(x) + g(x)) = 1 / (1 + g(x)/f(x))
  
      Thus define: u = 1 + g(x)/f(x)
      for which: du = [g'(x)h(x) - g(x)f'(x)] / h(x)^2 dx
      and solving integral 1/u du.
  
  
      #### References
  
        [1]: Jesse Davis and Mark Goadrich. The relationship between
             Precision-Recall and ROC curves. In _International Conference on
             Machine Learning_, 2006. http://dl.acm.org/citation.cfm?id=1143874
      """
      with tf.name_scope(name or "quantile_auc"):
          dtype = dtype_util.common_dtype(
              [q0, q1, n1, n0], dtype_hint=tf.float32
          )
          q1 = tf.convert_to_tensor(q1, dtype=dtype)
          q0 = tf.convert_to_tensor(q0, dtype=dtype)
          n1 = tf.convert_to_tensor(n1, dtype=dtype)
          n0 = tf.convert_to_tensor(n0, dtype=dtype)
  
          # Broadcast `q0` and `q1` to at least 2D. This makes `q1[i]` and `q0[j]` at
          # least 1D, allowing `tf.where` to operate on them.
          q1 = q1 + tf.zeros((1, 1), dtype=dtype)
          q0 = q0 + tf.zeros((1, 1), dtype=dtype)
  
          q1_shape = ps.shape(q1)
          batch_shape = q1_shape[:-1]
          n_q1 = q1_shape[-1]
          n_q0 = ps.shape(q0)[-1]
          static_batch_shape = q1.shape[:-1]
  
          n0 = tf.broadcast_to(n0, batch_shape)
          n1 = tf.broadcast_to(n1, batch_shape)
  
          def loop_body(auc_prev, p0_prev, p1_prev, i, j):
              """Body of the loop to integrate over the ROC/PR curves."""
  
              # We will align intervals from q0 and q1 by moving from end to beginning,
              # stopping with whatever boundary we encounter first. This boundary is
              # `b`. We find `a` by continuing the search from `b` to the left.
              b, i, j = _get_endpoint_b(i, j, q1, q0, batch_shape)
  
              # Get q1[i] and q1[i+1]
              ind = tf.where(tf.greater_equal(i, 0))
              q1_i = _get_q_slice(q1, i, ind, batch_shape=batch_shape)
              ip1 = tf.minimum(i + 1, n_q1 - 1)
              q1_ip1 = _get_q_slice(
                  q1, ip1, ind, batch_shape=batch_shape
              )
  
              # Get q0[j] and q0[j+1]
              ind = tf.where(tf.greater_equal(j, 0))
              q0_j = _get_q_slice(q0, j, ind, batch_shape=batch_shape)
              jp1 = tf.minimum(j + 1, n_q0 - 1)
              q0_jp1 = _get_q_slice(
                  q0, jp1, ind, batch_shape=batch_shape
              )
  
              a = _get_endpoint_a(i, j, q1_i, q0_j, batch_shape)
  
              # Calculate the proportion [a, b) represents of [q1[i], q1[i+1]).
              d1 = _get_interval_proportion(
                  i, n_q1, q1_i, q1_ip1, a, b, batch_shape
              )
  
              # Calculate the proportion [a, b) represents of [q1[i], q1[i+1]).
              d0 = _get_interval_proportion(
                  j, n_q0, q0_j, q0_jp1, a, b, batch_shape
              )
  
              # Notice that because we assumed within bucket values are distributed
              # uniformly, we end up with something which is identical to the
              # trapezoidal rule: definite_integral += (b - a) * (f(a) + f(b)) / 2.
              if curve == "ROC":
                  auc = auc_prev + d0 * (p1_prev + d1 / 2.0)
              else:
                  total_scaled_delta = n0 * d0 + n1 * d1
                  total_scaled_cdf_at_b = n0 * p0_prev + n1 * p1_prev
  
                  def get_auprc_update():
                      proportion = total_scaled_cdf_at_b / (
                          total_scaled_cdf_at_b + total_scaled_delta
                      )
                      definite_integral = (
                          n1 / tf.square(total_scaled_delta)
                      ) * (
                          d1 * total_scaled_delta
                          + tf.math.log(proportion)
                          * (
                              d1 * total_scaled_cdf_at_b
                              - p1_prev * total_scaled_delta
                          )
                      )
                      return d1 * definite_integral
  
                  # Values should be non-negative and we use > 0.0 rather than != 0.0 to
                  # catch possible numerical imprecision.
                  delta_gt_0 = tf.greater(total_scaled_delta, 0.0)
                  cdf_gt_0 = tf.greater(total_scaled_cdf_at_b, 0.0)
                  d1_gt_0 = tf.greater(d1, 0.0)
                  ind = tf.where(delta_gt_0 & cdf_gt_0 & d1_gt_0)
                  auc_update = tf.gather_nd(get_auprc_update(), ind)
                  auc = tf.tensor_scatter_nd_add(
                      auc_prev, ind, auc_update
                  )
  
              # TODO(jvdillon): In calculating AUROC and AUPRC, we probably should
              # resolve ties randomly, making the following eight states possible (where
              # e = 1 means we expected a positive trial and e = 0 means we expected a
              # negative trial):
              #
              #   P(y = 1 | pi(x) > delta)
              #   P(y = 1 | pi(x) = delta, e = 1) 0.5
              #   P(y = 1 | pi(x) = delta, e = 0) 0.5
              #   P(y = 1 | pi(x) < delta)
              #
              #   P(y = 0 | pi(x) > delta)
              #   P(y = 0 | pi(x) = delta, e = 1) 0.5
              #   P(y = 0 | pi(x) = delta, e = 0) 0.5
              #   P(y = 0 | pi(x) < delta)
              #
              # This makes the math a bit harder and its not clear this adds much,
              # especially when we're assuming piecewise uniformity.
  
              # Accumulate this mass (d1, d0) for the next iteration.
              p1 = p1_prev + d1
              p0 = p0_prev + d0
  
              return auc, p0, p1, i, j
  
          init_auc = tf.zeros(batch_shape, dtype=dtype)
          init_p0 = tf.zeros(batch_shape, dtype=dtype)
          init_p1 = tf.zeros(batch_shape, dtype=dtype)
          init_i = tf.zeros(batch_shape, dtype=tf.int64) + n_q1 - 1
          init_j = tf.zeros(batch_shape, dtype=tf.int64) + n_q0 - 1
  
          loop_cond = lambda auc, p0, p1, i, j: tf.reduce_all(  # pylint: disable=g-long-lambda
              tf.greater_equal(i, 0) | tf.greater_equal(j, 0)
          )
  
          init_vars = [init_auc, init_p0, init_p1, init_i, init_j]
          auc, _, _, _, _ = tf.while_loop(
              loop_cond,
              loop_body,
              init_vars,
              shape_invariants=[static_batch_shape] * 5,
          )
  
          return tf.squeeze(auc)
  
file_path: probability/tensorflow_probability/python/stats/ranking.py
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
label_desc: Use image node_id 1 for calling a global function with example usage: quantile_auc(q0, n0, q1, n1, curve, name) and returns: tf, auc,  p0,  p1,  i,  j, unknown