Repository: google/jax

Issue:
Abort trap in jax.numpy code
```python
import jax
import jaxlib

import jax.numpy as jnp

print(f"{jax.__version__=}")
print(f"{jaxlib.__version__=}")

x = jnp.arange(4)
y = jnp.arange(20)
z = jnp.ones(20, dtype='complex64')

print(jnp.interp(x, y, z, period=0.59))
```
output:
```
jax.__version__='0.3.14'
jaxlib.__version__='0.3.10'
2022-06-01 13:06:02.826868: F external/org_tensorflow/tensorflow/compiler/xla/service/hlo_evaluator.cc:137] unhandled direction for conversion to Comparison: LT
Abort trap: 6
```


Discussion:
I think it may have to do with a complex remainder operation
Certainly we should have rejected the `remainder` for complex types in several places before it made it to the compiler:
```
In [1]: import numpy as np, jax.numpy as jnp

In [2]: np.remainder(1+2j, 3)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Input In [2], in <cell line: 1>()
----> 1 np.remainder(1+2j, 3)

TypeError: ufunc 'remainder' not supported for the input types, and the inputs could not be safely coerced to any supported types according to the casting rule ''safe''

In [3]: jnp.remainder(1+2j, 3)
---------------------------------------------------------------------------
XlaRuntimeError                           Traceback (most recent call last)
Input In [3], in <cell line: 1>()
----> 1 jnp.remainder(1+2j, 3)

    [... skipping hidden 12 frame]

File ~/p/jax/jax/_src/dispatch.py:801, in backend_compile(backend, built_c, options)
    797 @profiler.annotate_function
    798 def backend_compile(backend, built_c, options):
    799   # we use a separate function call to ensure that XLA compilation appears
    800   # separately in Python profiling results
--> 801   return backend.compile(built_c, compile_options=options)

XlaRuntimeError: UNIMPLEMENTED: binary complex op 'remainder'
```
Here's a more compact repro:
```python
from jax import jit
import jax.numpy as jnp
from functools import partial

@partial(jit, static_argnums=1)
def f(x, a):
  return x % a

f(jnp.complex64(1), 0.5)
```
If the second argument is not annotated as static, we reject it before it makes it to the compiler.

Resolving patch:
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -1552,6 +1552,7 @@ def _nary_lower_mhlo(op: Callable, ctx,
 _num = _int | _float | _complex
 _any = _int | _float | _complex | _bool
 _bool_or_int = _int | _bool
+_ordered = _int | _float | _bool
 
 neg_p = standard_unop(_num, 'neg')
 ad.deflinear2(neg_p, lambda t, operand: [neg(t)])
@@ -2106,7 +2107,7 @@ def _div_transpose_rule(cotangent, x, y):
 ad.primitive_transposes[div_p] = _div_transpose_rule
 mlir.register_lowering(div_p, partial(_nary_lower_mhlo, mhlo.DivOp))
 
-rem_p = standard_naryop([_num, _num], 'rem')
+rem_p = standard_naryop([_int | _float, _int | _float], 'rem')
 ad.defjvp(
     rem_p,
     lambda g, x, y: _maybe_broadcast(broadcast_shapes(np.shape(x), np.shape(y)), g),
@@ -2171,19 +2172,19 @@ def _compare_lower_mhlo(direction: str, ctx, x, y):
 ad.defjvp_zero(ne_p)
 mlir.register_lowering(ne_p, partial(_compare_lower_mhlo, "NE"))
 
-ge_p = naryop(_fixed_dtype(np.bool_), [_any, _any], 'ge')
+ge_p = naryop(_fixed_dtype(np.bool_), [_ordered, _ordered], 'ge')
 ad.defjvp_zero(ge_p)
 mlir.register_lowering(ge_p, partial(_compare_lower_mhlo, "GE"))
 
-gt_p = naryop(_fixed_dtype(np.bool_), [_any, _any], 'gt')
+gt_p = naryop(_fixed_dtype(np.bool_), [_ordered, _ordered], 'gt')
 ad.defjvp_zero(gt_p)
 mlir.register_lowering(gt_p, partial(_compare_lower_mhlo, "GT"))
 
-le_p = naryop(_fixed_dtype(np.bool_), [_any, _any], 'le')
+le_p = naryop(_fixed_dtype(np.bool_), [_ordered, _ordered], 'le')
 ad.defjvp_zero(le_p)
 mlir.register_lowering(le_p, partial(_compare_lower_mhlo, "LE"))
 
-lt_p = naryop(_fixed_dtype(np.bool_), [_any, _any], 'lt')
+lt_p = naryop(_fixed_dtype(np.bool_), [_ordered, _ordered], 'lt')
 ad.defjvp_zero(lt_p)
 mlir.register_lowering(lt_p, partial(_compare_lower_mhlo, "LT"))
 


Test patch:
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -2533,6 +2533,11 @@ def test_reduction_with_repeated_axes_error(self):
     with self.assertRaisesRegex(ValueError, "duplicate value in 'axes' .*"):
       lax.reduce(np.arange(3), 0, lax.add, (0, 0))
 
+  @parameterized.parameters([lax.rem, lax.lt, lax.gt, lax.ge, lax.le])
+  def test_ops_do_not_accept_complex_dtypes(self, op):
+    with self.assertRaisesRegex(TypeError, ".*does not accept dtype complex.*"):
+      op(2+3j, 4+5j)
+
   def test_population_count_booleans_not_supported(self):
     # https://github.com/google/jax/issues/3886
     msg = "population_count does not accept dtype bool"