Changing implementation to one that works in all modes.

PiperOrigin-RevId: 463143711

Author: Googler

tensorflow_probability/python/stats/BUILD

@@ -33,7 +33,6 @@ multi_substrate_py_library(
     name = "stats",
     srcs = ["__init__.py"],
     jax_omit_deps = [
-        ":kendalls_tau",
         ":ranking",
     ],
     deps = [

tensorflow_probability/python/stats/__init__.py

@@ -20,9 +20,7 @@
 from tensorflow_probability.python.stats.calibration import brier_score
 from tensorflow_probability.python.stats.calibration import expected_calibration_error
 from tensorflow_probability.python.stats.calibration import expected_calibration_error_quantiles
-from tensorflow_probability.python.stats.kendalls_tau import iterative_mergesort
 from tensorflow_probability.python.stats.kendalls_tau import kendalls_tau
-from tensorflow_probability.python.stats.kendalls_tau import lexicographical_indirect_sort
 from tensorflow_probability.python.stats.leave_one_out import log_loomean_exp
 from tensorflow_probability.python.stats.leave_one_out import log_loosum_exp
 from tensorflow_probability.python.stats.leave_one_out import log_soomean_exp
@@ -65,9 +63,7 @@
     'expected_calibration_error_quantiles',
     'find_bins',
     'histogram',
-    'iterative_mergesort',
     'kendalls_tau',
-    'lexicographical_indirect_sort',
     'log_average_probs',
     'log_loomean_exp',
     'log_loosum_exp',

tensorflow_probability/python/stats/kendalls_tau.py

@@ -21,149 +21,42 @@
 from tensorflow_probability.python.internal import prefer_static as ps
 from tensorflow_probability.python.internal import tensorshape_util
 
-__all__ = ['iterative_mergesort', 'kendalls_tau']
+__all__ = ['kendalls_tau']
 
 
-def iterative_mergesort(y, permutation, name=None):
-  """Non-recusive mergesort that counts exchanges.
+def _tril_indices(n):
+  """Emulate np.tril_indices(n, k=-1).
+
+  This method ensures static shapes throughout (ie, XLA compilable).
+  This method only works for n <= 30000.
 
   Args:
-    y: a `Tensor` of shape `[n]` containing values to be sorted.
-    permutation: `Tensor` of shape `[n]` with original ordering.
-    name: Optional Python `str` name for ops created by this method.
-      Default value: `None` (i.e., 'iterative_mergesort').
+    n: number elements to generate all pairs
 
   Returns:
-    exchanges: `int32` scalar that counts the number of exchanges required to
-      produce a sorted permutation
-    permutation: and a `tf.int32` Tensor that contains the ordering of y values
-      that are sorted.
+    A [2, n * (n - 1) / 2] vector of all combinations of range(n).
   """
+  n = tf.convert_to_tensor(n, dtype_hint=tf.int32)
+  # Number of lower triangular entries in an nxn matrix
+  m = (n - 1) * n / 2
+  r = tf.cast(tf.range(m), dtype=tf.float64)
 
-  with tf.name_scope(name or 'iterative_mergesort'):
-    y = tf.convert_to_tensor(y, name='y')
-    permutation = tf.convert_to_tensor(
-        permutation, name='permutation', dtype=tf.int32)
-    shape = permutation.shape
-    tensorshape_util.assert_is_compatible_with(y.shape, shape)
-    n = ps.size(y)
-
-    def outer_body(k, exchanges, permutation):
-      # The outer body progressively merges lists as k grows by powers of 2,
-      # tracking the total swaps required in exchanges as the new permutation is
-      # built in place.
-      y_ordered = tf.gather(y, permutation)
-
-      def middle_body(left, exchanges, permutation):
-        # the middle body advances through the sublists of size k, advancing
-        # the left edge until the end of the input is reached.
-        right = left + k
-        end = tf.minimum(right + k, n)
-
-        # See explanation here
-        # https://www.geeksforgeeks.org/counting-inversions/.
-
-        def inner_body(i, j, x, np, p):
-          # The [left, right) and [right, end) lists are merged sorted, with
-          # i and j tracking the advance through each range. x records the
-          # number of order (bubble-sort equivalent) swaps that are happening
-          # with each insertion, and np represents the size of the output
-          # permutation that's been filled in using the p tensor.
-          y_less = y_ordered[i] <= y_ordered[j]
-          element = tf.where(y_less, [permutation[i]], [permutation[j]])
-          new_p = tf.concat([p[0:np], element, p[np + 1:n]], axis=0)
-          tensorshape_util.set_shape(new_p, p.shape)
-          return (tf.where(y_less, i + 1, i), tf.where(y_less, j, j + 1),
-                  tf.where(y_less, x, x + right - i), np + 1, new_p)
-
-        i_j_x_np_p = (left, right, exchanges, 0, tf.zeros([n], dtype=tf.int32))
-        (i, j, exchanges, np, p) = tf.while_loop(
-            cond=lambda i, j, x, np, p: tf.math.logical_and(i < right, j < end),
-            body=inner_body,
-            loop_vars=i_j_x_np_p)
-        permutation = tf.concat([
-            permutation[0:left], p[0:np], permutation[i:right],
-            permutation[j:end], permutation[end:n]
-        ],
-                                axis=0)
-        tensorshape_util.set_shape(permutation, shape)
-        return left + 2 * k, exchanges, permutation
-
-      _, exchanges, permutation = tf.while_loop(
-          cond=lambda left, exchanges, permutation: left < n - k,
-          body=middle_body,
-          loop_vars=(0, exchanges, permutation))
-      k *= 2
-      return k, exchanges, permutation
-
-    _, exchanges, permutation = tf.while_loop(
-        cond=lambda k, exchanges, permutation: k < n,
-        body=outer_body,
-        loop_vars=(1, 0, permutation))
-    return exchanges, permutation
-
-
-def lexicographical_indirect_sort(primary, secondary, name=None):
-  """Sorts by primary, then by secondary returning the indices.
+  # From Sloane: https://oeis.org/A002024 "k appears k times"
+  # e.g., [1, 2, 2, 3, 3, 3, 4, 4, 4, 4, ...]
+  e = tf.math.floor(tf.math.sqrt(2 * (r + 1)) + .5)
 
-  Args:
-    primary: a `Tensor` of shape `[n]` containing the primary sort key. the
-      primary sort key value.
-    secondary: a `Tensor` of shape `[n]` containing the secondary sort key to be
-      used when the primary keys are identical.
-    name: Optional Python `str` name for ops created by this method.
-      Default value: `None` (i.e., 'lexicographical_indirect_sort').
+  # From Sloane: https://oeis.org/A002262 "Triangle read by rows"
+  # e.g., [0, 0, 1, 0, 1, 2, 0, 1, 2, 3, ...]
+  f = tf.math.floor(tf.math.sqrt(2 * r + .25) - .5)
+  g = r - f * (f + 1) / 2
 
-  Returns:
-    lexicographic: A permutation of range(n) that provides the sorted primary,
-      then secondary values.
-  """
-  with tf.name_scope(name or 'lexicographical_indirect_sort'):
-    n = ps.size0(primary)
-    permutation = tf.argsort(primary)
-    # scan for ties, and for each range of ties do a argsort on
-    # the secondary value. (TF has no lexicographical sorting, although
-    # jax can sort complex number lexicographically. Hmm.)
-    primary_ordered = tf.gather(primary, permutation)
-
-    def body(left, right, lexicographic):
-      # We make a single pass through the list using right and left, where right
-      # advances and left chases it looking for spans that are equal in their
-      # primary key to then institute a sort on the secondary key.
-      not_equal = tf.not_equal(primary_ordered[left], primary_ordered[right])
-
-      def secondary_sort():
-        x = tf.concat([
-            lexicographic[0:left],
-            tf.gather(permutation[left:right],
-                      tf.argsort(tf.gather(secondary,
-                                           permutation[left:right]))),
-            lexicographic[right:n],
-        ],
-                      axis=0)
-        tensorshape_util.set_shape(x, [n])
-        return x
-
-      return (tf.where(not_equal, right, left), right + 1,
-              tf.cond(not_equal, secondary_sort, lambda: lexicographic))
-
-    left, _, lexicographic = tf.while_loop(
-        cond=lambda left, right, lexicographic: right < n,
-        body=body,
-        loop_vars=(0, 0, tf.zeros_like(permutation, dtype=tf.int32)))
-    return tf.concat([
-        lexicographic[0:left],
-        tf.gather(permutation[left:n],
-                  tf.argsort(tf.gather(secondary, permutation[left:n])))
-    ],
-                     axis=0)
+  return tf.cast(tf.stack([e, g]), dtype=tf.int32)
 
 
 def kendalls_tau(y_true, y_pred, name=None):
   """Computes Kendall's Tau for two ordered lists.
 
-  Kendall's Tau measures the correlation between ordinal rankings. This
-  implementation is similar to the one used in scipy.stats.kendalltau.
+  Kendall's Tau measures the correlation between ordinal rankings.
   The provided values may be of any type that is sortable, with the
   argsort indices indicating the true or proposed ordinal sequence.
 
@@ -189,62 +82,21 @@ def kendalls_tau(y_true, y_pred, name=None):
             ps.size(y_true), 1, 'Ordering requires at least 2 elements.')
     ]
     with tf.control_dependencies(assertions):
-      lexa = lexicographical_indirect_sort(y_true, y_pred)
-
-    # See A Computer Method for Calculating Kendall's Tau with Ungrouped Data
-    # by William Night, Journal of the American Statistical Association,
-    # Jun., 1966, Vol. 61, No. 314, Part 1 (Jun., 1966), pp. 436-439
-    # for notation https://www.jstor.org/stable/2282833
-
-    def jointly_tied_pairs_body(first, t, i):
-      not_equal = tf.math.logical_or(
-          tf.not_equal(y_true[lexa[first]], y_true[lexa[i]]),
-          tf.not_equal(y_pred[lexa[first]], y_pred[lexa[i]]))
-      return (tf.where(not_equal, i, first),
-              tf.where(not_equal, t + ((i - first) * (i - first - 1)) // 2,
-                       t), i + 1)
-
-    n = ps.size0(y_true)
-    first, t, _ = tf.while_loop(
-        cond=lambda first, t, i: i < n,
-        body=jointly_tied_pairs_body,
-        loop_vars=(0, 0, 1))
-    t += ((n - first) * (n - first - 1)) // 2
-
-    def ties_y_true_body(first, v, i):
-      not_equal = tf.not_equal(y_true[lexa[first]], y_true[lexa[i]])
-      return (tf.where(not_equal, i, first),
-              tf.where(not_equal, v + ((i - first) * (i - first - 1)) // 2,
-                       v), i + 1)
-
-    first, v, _ = tf.while_loop(
-        cond=lambda first, v, i: i < n,
-        body=ties_y_true_body,
-        loop_vars=(0, 0, 1))
-    v += ((n - first) * (n - first - 1)) // 2
-
-    # count exchanges
-    exchanges, newperm = iterative_mergesort(y_pred, lexa)
-
-    def ties_in_y_pred_body(first, u, i):
-      not_equal = tf.not_equal(y_pred[newperm[first]], y_pred[newperm[i]])
-      return (tf.where(not_equal, i, first),
-              tf.where(not_equal, u + ((i - first) * (i - first - 1)) // 2,
-                       u), i + 1)
-
-    first, u, _ = tf.while_loop(
-        cond=lambda first, u, i: i < n,
-        body=ties_in_y_pred_body,
-        loop_vars=(0, 0, 1))
-    u += ((n - first) * (n - first - 1)) // 2
-    n0 = (n * (n - 1)) // 2
-    assertions = [
-        assert_util.assert_less(v, tf.cast(n0, tf.int32),
-                                'All ranks are ties for y_true.'),
-        assert_util.assert_less(u, tf.cast(n0, tf.int32),
-                                'All ranks are ties for y_pred.')
-    ]
-    with tf.control_dependencies(assertions):
-      return (tf.cast(n0 - (u + v - t), tf.float32) -
-              2.0 * tf.cast(exchanges, tf.float32)) / tf.math.sqrt(
-                  tf.cast(n0 - v, tf.float32) * tf.cast(n0 - u, tf.float32))
+      n = ps.size0(y_true)
+      indices = _tril_indices(n)
+      dxij = tf.sign(
+          tf.gather(y_true, indices[0]) - tf.gather(y_true, indices[1]))
+      dyij = tf.sign(
+          tf.gather(y_pred, indices[0]) - tf.gather(y_pred, indices[1]))
+      # s is sum of concordant pairs minus discordant pairs.
+      s = tf.cast(tf.math.reduce_sum(dxij * dyij), tf.float32)
+      # t is the number of y_true pairs that are not ties.
+      t = tf.math.count_nonzero(dxij, dtype=tf.float32)
+      # u is the number of y_pred pairs that are not ties.
+      u = tf.math.count_nonzero(dyij, dtype=tf.float32)
+      assertions = [
+          assert_util.assert_positive(t, 'All ranks are ties for y_true.'),
+          assert_util.assert_positive(u, 'All ranks are ties for y_pred.')
+      ]
+      with tf.control_dependencies(assertions):
+        return s / tf.math.sqrt(t * u)

tensorflow_probability/python/stats/kendalls_tau_test.py

@@ -26,17 +26,6 @@
 @test_util.test_all_tf_execution_regimes
 class KendallsTauTest(test_util.TestCase):
 
-  def test_iterative_mergesort(self):
-    values = [7, 3, 9, 0, -6, 12, 54, 3, -6, 88, 1412]
-    array = tf.constant(values, tf.int32)
-    iperm = tf.range(len(values), dtype=tf.int32)
-    exchanges, perm = tfp.stats.iterative_mergesort(array, iperm)
-    expected = sorted(values)
-    self.assertAllEqual(expected, tf.gather(array, perm))
-    ordered, _ = tfp.stats.iterative_mergesort(array, perm)
-    self.assertAllEqual(ordered, 0)
-    self.assertAllEqual(exchanges, 19)
-
   def test_kendall_tau(self):
     x1 = [12, 2, 1, 12, 2]
     x2 = [1, 4, 7, 1, 0]
@@ -46,14 +35,6 @@ def test_kendall_tau(self):
             tf.constant(x1, tf.float32), tf.constant(x2, tf.float32)))
     self.assertAllClose(expected, res, atol=1e-5)
 
-  def test_lexicographical_sort(self):
-    primary = [12, 2, 1, 12, 2]
-    secondary = [1, 4, 7, 1, 0]
-    expected = [2, 4, 1, 0, 3]  # Assumes stable sort.
-    res = self.evaluate(
-        tfp.stats.lexicographical_indirect_sort(primary, secondary))
-    self.assertAllEqual(expected, res)
-
   def test_kendall_tau_float(self):
     x1 = [0.12, 0.02, 0.01, 0.12, 0.02]
     x2 = [0.1, 0.4, 0.7, 0.1, 0.0]