Repository: google/jax

Issue:
RuntimeError: Unimplemented: complex comparison 'LT'
I'm getting a weird internal JAX error with this script:
```python
import jax.numpy as jp

def lqr_continuous_time_infinite_horizon(A, B, Q, R, N):
  # Take the last dimension, in case we try to do some kind of broadcasting
  # thing in the future.
  x_dim = A.shape[-1]

  # See https://en.wikipedia.org/wiki/Linear%E2%80%93quadratic_regulator#Infinite-horizon,_continuous-time_LQR.
  A1 = A - B @ jp.linalg.solve(R, N.T)
  Q1 = Q - N @ jp.linalg.solve(R, N.T)

  # See https://en.wikipedia.org/wiki/Algebraic_Riccati_equation#Solution.
  H = jp.block([[A1, -B @ jp.linalg.solve(R, B.T)], [-Q1, -A1]])
  eigvals, eigvectors = jp.linalg.eig(H)
  argsort = jp.argsort(eigvals)
  ix = argsort[:x_dim]
  U = eigvectors[:, ix]
  P = U[x_dim:, :] @ jp.linalg.inv(U[:x_dim, :])

  K = jp.linalg.solve(R, (B.T @ P + N.T))
  return K, P, eigvals[ix]

def _test_lqr(n):
  import control
  from jax.tree_util import tree_multimap

  A = jp.eye(n)
  B = jp.eye(n)
  Q = jp.eye(n)
  R = jp.eye(n)
  N = jp.zeros((n, n))

  actual = lqr_continuous_time_infinite_horizon(A, B, Q, R, N)
  expected = control.lqr(A, B, Q, R, N)
  assert tree_multimap(jp.allclose, actual, expected)

if __name__ == "__main__":
  _test_lqr(2)
```
I'm getting:
```
❯ pipenv run python -m research.lqr
/Users/skainswo/dev/jax/jax/lib/xla_bridge.py:116: UserWarning: No GPU/TPU found, falling back to CPU.
  warnings.warn('No GPU/TPU found, falling back to CPU.')
Traceback (most recent call last):
  File "/usr/local/Cellar/python/3.7.6_1/Frameworks/Python.framework/Versions/3.7/lib/python3.7/runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "/usr/local/Cellar/python/3.7.6_1/Frameworks/Python.framework/Versions/3.7/lib/python3.7/runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "/Users/skainswo/dev/research/research/lqr.py", line 50, in <module>
    _test_lqr1(2)
  File "/Users/skainswo/dev/research/research/lqr.py", line 45, in _test_lqr1
    actual = lqr_continuous_time_infinite_horizon(A, B, Q, R, N)
  File "/Users/skainswo/dev/research/research/lqr.py", line 26, in lqr_continuous_time_infinite_horizon
    argsort = jp.argsort(eigvals)
  File "/Users/skainswo/dev/jax/jax/numpy/lax_numpy.py", line 2886, in argsort
    _, perm = lax.sort_key_val(a, iota, dimension=axis)
  File "/Users/skainswo/dev/jax/jax/lax/lax.py", line 1190, in sort_key_val
    result = sort_key_val_p.bind(keys, values, dimension=dimension)
  File "/Users/skainswo/dev/jax/jax/core.py", line 211, in bind
    return self.impl(*args, **kwargs)
  File "/Users/skainswo/dev/jax/jax/interpreters/xla.py", line 217, in apply_primitive
    compiled_fun = xla_primitive_callable(prim, *map(arg_spec, args), **params)
  File "/Users/skainswo/dev/jax/jax/interpreters/xla.py", line 254, in xla_primitive_callable
    compiled = backend.compile(built_c, compile_options=options)
RuntimeError: Unimplemented: complex comparison 'LT'
```
I'm guessing that this has something to do with the fact that I'm getting complex eigenvalues out, but the error message is pretty confusing...


Discussion:
Yes, this is certainly a bad error.

I'm curious if you actually want the `np.argsort` behavior on complex numbers (lexicographic order) or just a better error... either way, I guess we'll implement the `np.argsort` behavior.
@hawkinsp Yeah, that's a great question. I think I'd personally prefer a better error, but only because it's not immediately clear to me what the right ordering should be.

Resolving patch:
diff --git a/jax/lax/lax.py b/jax/lax/lax.py
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -4289,6 +4289,11 @@ def _select_and_gather_add_shape_rule(
   64: onp.uint64,
 }
 
+_INT_DTYPES = {
+  16: onp.int16,
+  32: onp.int32,
+  64: onp.int64,
+}
 
 def _select_and_gather_add_translation(
     c, tangents, operand, *, select_prim, window_dimensions, window_strides,
@@ -4563,8 +4568,69 @@ def _sort_abstract_eval(*args, **kwargs):
     raise TypeError(f"Arguments to sort must have equal shapes, got: {shapes}")
   return args
 
+
+def _float_to_int_for_sort(x):
+  # Switch from a floating point value to a integer value in such a way that
+  # when using the integer value to compare, we get the same result for normal
+  # values, and -nan is treated as the smallest value, and nan is treated as
+  # the largest value.
+  # If f is a float, and
+  # x = bit_cast<int32>(f);
+  # y = x < 0 ? int32_max - x : x;
+  # then y is ordered as an int32 such that finite values have the obvious
+  # order, -0 is ordered before 0, and -NaN and NaN appear at the beginning
+  # and end of the ordering.
+  # Note that in order to avoid -x to overflow, we calculate
+  # int32_max - x as unsigned, and then convert back to signed.
+  if x.dtype == dtypes.bfloat16:
+    x = convert_element_type(x, onp.float32)
+  nbits = onp.finfo(x).bits
+  signed_dtype = _INT_DTYPES[nbits]
+  unsigned_dtype = _UINT_DTYPES[nbits]
+
+  signed = bitcast_convert_type(x, signed_dtype)
+  unsigned = bitcast_convert_type(x, unsigned_dtype)
+  flipped = bitcast_convert_type(
+    sub(unsigned_dtype(onp.iinfo(signed_dtype).max), unsigned), signed_dtype)
+  return select(lt(signed, _zero(signed)), flipped, signed)
+
+# Default comparator that sorts the operands only on their first arguments.
+# For floating point types, a total order is created where
+# -NaN < -infinity < ... < -0 < 0 < ... < infinity < NaN.
+# For complex types, the (real, imag) pairs are sorted lexicographically
+# (following NumPy's semantics).
+# This code adds complex-number support to the algorithm from:
+# https://github.com/tensorflow/tensorflow/blob/ba43780830f09da72081fe5061c436f1c6203a92/tensorflow/compiler/xla/client/lib/comparators.h#L33
+def _sort_lt_comparator(*operands):
+  assert len(operands) >= 2 and len(operands) % 2 == 0, operands
+  x, y = operands[:2]
+  assert x.dtype == y.dtype, (x.dtype, y.dtype)
+  if onp.issubdtype(x.dtype, onp.complexfloating):
+    x_keys = [_float_to_int_for_sort(real(x)), _float_to_int_for_sort(imag(x))]
+    y_keys = [_float_to_int_for_sort(real(y)), _float_to_int_for_sort(imag(y))]
+  elif onp.issubdtype(x.dtype, onp.floating):
+    x_keys = [_float_to_int_for_sort(x)]
+    y_keys = [_float_to_int_for_sort(y)]
+  else:
+    x_keys = [x]
+    y_keys = [y]
+
+  p = None
+  for xk, yk in zip(x_keys[::-1], y_keys[::-1]):
+    p = (bitwise_or(lt(xk, yk), bitwise_and(eq(xk, yk), p)) if p is not None
+         else lt(xk, yk))
+  return p
+
 def _sort_translation_rule(c, *operands, dimension):
-  out = xops.Sort(c, operands, dimension=dimension, is_stable=True)
+  types = [c.get_shape(x).xla_element_type() for x in operands]
+  subc = xla_bridge.make_computation_builder("sort_lt_comparator")
+  params = [xb.parameter(subc, 2 * i + j, xc.Shape.array_shape(typ, ()))
+            for i, typ in enumerate(types) for j in range(2)]
+  result = xla.lower_fun(_sort_lt_comparator,
+                         multiple_results=False)(subc, *params)
+  comparator = subc.build(result)
+  out = xops.Sort(c, operands, dimension=dimension, is_stable=True,
+                  comparator=comparator)
   return out if len(operands) != 1 else xops.Tuple(c, [out])
 
 def _sort_jvp(primals, tangents, *, dimension):


Test patch:
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -1302,13 +1302,17 @@ def testCumulativeReduce(self, op, onp_op, shape, dtype, axis, rng_factory):
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_shape={}_axis={}".format(
           jtu.format_shape_dtype_string(shape, dtype), axis),
-       "rng_factory": rng_factory, "shape": shape, "dtype": dtype, "axis": axis}
-      for dtype in [onp.float32, onp.int32, onp.uint32]
+       "shape": shape, "dtype": dtype, "axis": axis}
+      for dtype in all_dtypes
       for shape in [(5,), (5, 7)]
-      for axis in [-1, len(shape) - 1]
-      for rng_factory in [jtu.rand_default]))
-  def testSort(self, shape, dtype, axis, rng_factory):
-    rng = rng_factory(self.rng())
+      for axis in [-1, len(shape) - 1]))
+  def testSort(self, shape, dtype, axis):
+    # TODO(b/141131288): enable complex-valued sorts on TPU.
+    if (onp.issubdtype(dtype, onp.complexfloating) and (
+        (jtu.device_under_test() == "cpu" and jax.lib.version <= (0, 1, 47)) or
+         jtu.device_under_test() == "tpu")):
+      raise SkipTest("Complex-valued sort not implemented")
+    rng = jtu.rand_default(self.rng())
     args_maker = lambda: [rng(shape, dtype)]
     fun = lambda x: lax.sort(x, dimension=axis)
     self._CompileAndCheck(fun, args_maker, check_dtypes=True)
@@ -1316,13 +1320,17 @@ def testSort(self, shape, dtype, axis, rng_factory):
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_shape={}_axis={}".format(
           jtu.format_shape_dtype_string(shape, dtype), axis),
-       "rng_factory": rng_factory, "shape": shape, "dtype": dtype, "axis": axis}
-      for dtype in [onp.float32, onp.int32, onp.uint32]
+        "shape": shape, "dtype": dtype, "axis": axis}
+      for dtype in all_dtypes
       for shape in [(5,), (5, 7)]
-      for axis in [-1, len(shape) - 1]
-      for rng_factory in [jtu.rand_default]))
-  def testSortAgainstNumpy(self, shape, dtype, axis, rng_factory):
-    rng = rng_factory(self.rng())
+      for axis in [-1, len(shape) - 1]))
+  def testSortAgainstNumpy(self, shape, dtype, axis):
+    # TODO(b/141131288): enable complex-valued sorts on TPU.
+    if (onp.issubdtype(dtype, onp.complexfloating) and (
+        (jtu.device_under_test() == "cpu" and jax.lib.version <= (0, 1, 47)) or
+         jtu.device_under_test() == "tpu")):
+      raise SkipTest("Complex-valued sort not implemented")
+    rng = jtu.rand_default(self.rng())
     args_maker = lambda: [rng(shape, dtype)]
     op = lambda x: lax.sort(x, dimension=axis)
     numpy_op = lambda x: lax_reference.sort(x, axis)
@@ -1333,15 +1341,19 @@ def testSortAgainstNumpy(self, shape, dtype, axis, rng_factory):
           jtu.format_shape_dtype_string(shape, key_dtype),
           jtu.format_shape_dtype_string(shape, val_dtype),
           axis),
-       "rng_factory": rng_factory, "shape": shape,
-       "key_dtype": key_dtype, "val_dtype": val_dtype, "axis": axis}
-      for key_dtype in [onp.float32, onp.int32, onp.uint32]
+       "shape": shape, "key_dtype": key_dtype, "val_dtype": val_dtype,
+       "axis": axis}
+      for key_dtype in float_dtypes + complex_dtypes + int_dtypes + uint_dtypes
       for val_dtype in [onp.float32, onp.int32, onp.uint32]
       for shape in [(3,), (5, 3)]
-      for axis in [-1, len(shape) - 1]
-      for rng_factory in [jtu.rand_default]))
-  def testSortKeyVal(self, shape, key_dtype, val_dtype, axis, rng_factory):
-    rng = rng_factory(self.rng())
+      for axis in [-1, len(shape) - 1]))
+  def testSortKeyVal(self, shape, key_dtype, val_dtype, axis):
+    # TODO(b/141131288): enable complex-valued sorts on TPU.
+    if (onp.issubdtype(key_dtype, onp.complexfloating) and (
+        (jtu.device_under_test() == "cpu" and jax.lib.version <= (0, 1, 47)) or
+         jtu.device_under_test() == "tpu")):
+      raise SkipTest("Complex-valued sort not implemented")
+    rng = jtu.rand_default(self.rng())
     # This test relies on the property that wherever keys are tied, values are
     # too, since we don't guarantee the same ordering of values with equal keys.
     # To avoid that case, we generate unique keys (globally in the key array).
@@ -1359,15 +1371,19 @@ def args_maker():
           jtu.format_shape_dtype_string(shape, key_dtype),
           jtu.format_shape_dtype_string(shape, val_dtype),
           axis),
-       "rng_factory": rng_factory, "shape": shape,
-       "key_dtype": key_dtype, "val_dtype": val_dtype, "axis": axis}
-      for key_dtype in [onp.float32, onp.int32, onp.uint32]
+       "shape": shape, "key_dtype": key_dtype, "val_dtype": val_dtype,
+       "axis": axis}
+      for key_dtype in float_dtypes + complex_dtypes + int_dtypes + uint_dtypes
       for val_dtype in [onp.float32, onp.int32, onp.uint32]
       for shape in [(3,), (5, 3)]
-      for axis in [-1, len(shape) - 1]
-      for rng_factory in [jtu.rand_default]))
-  def testSortKeyValAgainstNumpy(self, shape, key_dtype, val_dtype, axis, rng_factory):
-    rng = rng_factory(self.rng())
+      for axis in [-1, len(shape) - 1]))
+  def testSortKeyValAgainstNumpy(self, shape, key_dtype, val_dtype, axis):
+    # TODO(b/141131288): enable complex-valued sorts on TPU.
+    if (onp.issubdtype(key_dtype, onp.complexfloating) and (
+        (jtu.device_under_test() == "cpu" and jax.lib.version <= (0, 1, 47)) or
+         jtu.device_under_test() == "tpu")):
+      raise SkipTest("Complex-valued sort not implemented")
+    rng = jtu.rand_default(self.rng())
     # This test relies on the property that wherever keys are tied, values are
     # too, since we don't guarantee the same ordering of values with equal keys.
     # To avoid that case, we generate unique keys (globally in the key array).