Repository: google/jax

Issue:
Type specification/inferrence now fails in @jax.jit
Please:

- [x] Check for duplicate issues.
- [x] Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:

The following [test case](https://github.com/ska-sa/codex-africanus/blob/master/africanus/rime/jax/tests/test_jax_phase_delay.py) has worked for a while, but has recently started failing. It seems that there isn't a constant handler for `complex64` anymore? The following reproducer demonstrates the issue, but replacing `out_dtype.type(1j)` with `1j` fixes the problem.

This works on jax 0.2.26 and jaxlib 0.1.75, but fails on jax 0.2.27 and 0.1.76.

```python
import numpy as np
import jax
import jax.numpy as jnp

@jax.jit
def phase_delay(lm, uvw, frequency):
    out_dtype = jnp.result_type(lm, uvw, frequency, np.complex64)

    one = lm.dtype.type(1.0)
    neg_two_pi_over_c = lm.dtype.type(-2*np.pi/3e8)

    l = lm[:, 0, None, None]  # noqa
    m = lm[:, 1, None, None]

    u = uvw[None, :, 0, None]
    v = uvw[None, :, 1, None]
    w = uvw[None, :, 2, None]

    n = jnp.sqrt(one - l**2 - m**2) - one

    real_phase = (neg_two_pi_over_c *
                  (l * u + m * v + n * w) *
                  frequency[None, None, :])
 
    # replacing out_dtype.type(1j) with 1j fixes this problem
    return jnp.exp(out_dtype.type(1j)*real_phase)

if __name__ == "__main__":
    uvw = np.random.random(size=(100, 3)).astype(np.float32)
    lm = np.random.random(size=(10, 2)).astype(np.float32)*0.001
    frequency = np.linspace(.856e9, .856e9*2, 64).astype(np.float32)

    complex_phase = phase_delay(lm, uvw, frequency)
```


- [x] #9390

```
$ python test_complex_constant_fail.py 
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
Traceback (most recent call last):
  File "test_complex_constant_fail.py", line 32, in <module>
    complex_phase = phase_delay(lm, uvw, frequency)
  File "/home/sperkins/venv/afr/lib/python3.7/site-packages/jax/_src/traceback_util.py", line 165, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/home/sperkins/venv/afr/lib/python3.7/site-packages/jax/_src/api.py", line 429, in cache_miss
    donated_invars=donated_invars, inline=inline)
  File "/home/sperkins/venv/afr/lib/python3.7/site-packages/jax/core.py", line 1671, in bind
    return call_bind(self, fun, *args, **params)
  File "/home/sperkins/venv/afr/lib/python3.7/site-packages/jax/core.py", line 1683, in call_bind
    outs = top_trace.process_call(primitive, fun, tracers, params)
  File "/home/sperkins/venv/afr/lib/python3.7/site-packages/jax/core.py", line 596, in process_call
    return primitive.impl(f, *tracers, **params)
  File "/home/sperkins/venv/afr/lib/python3.7/site-packages/jax/_src/dispatch.py", line 143, in _xla_call_impl
    *unsafe_map(arg_spec, args))
  File "/home/sperkins/venv/afr/lib/python3.7/site-packages/jax/linear_util.py", line 272, in memoized_fun
    ans = call(fun, *args)
  File "/home/sperkins/venv/afr/lib/python3.7/site-packages/jax/_src/dispatch.py", line 170, in _xla_callable_uncached
    *arg_specs).compile().unsafe_call
  File "/home/sperkins/venv/afr/lib/python3.7/site-packages/jax/_src/profiler.py", line 206, in wrapper
    return func(*args, **kwargs)
  File "/home/sperkins/venv/afr/lib/python3.7/site-packages/jax/_src/dispatch.py", line 260, in lower_xla_callable
    donated_invars)
  File "/home/sperkins/venv/afr/lib/python3.7/site-packages/jax/interpreters/mlir.py", line 403, in lower_jaxpr_to_module
    input_output_aliases=input_output_aliases)
  File "/home/sperkins/venv/afr/lib/python3.7/site-packages/jax/interpreters/mlir.py", line 541, in lower_jaxpr_to_fun
    *args)
  File "/home/sperkins/venv/afr/lib/python3.7/site-packages/jax/interpreters/mlir.py", line 606, in jaxpr_subcomp
    in_nodes = map(read, eqn.invars)
  File "/home/sperkins/venv/afr/lib/python3.7/site-packages/jax/_src/util.py", line 44, in safe_map
    return list(map(f, *args))
  File "/home/sperkins/venv/afr/lib/python3.7/site-packages/jax/interpreters/mlir.py", line 583, in read
    return ir_constants(v.val, canonicalize_types=True)
  File "/home/sperkins/venv/afr/lib/python3.7/site-packages/jax/interpreters/mlir.py", line 171, in ir_constants
    raise TypeError("No constant handler for type: {}".format(type(val)))
jax._src.traceback_util.UnfilteredStackTrace: TypeError: No constant handler for type: <class 'numpy.complex64'>

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "test_complex_constant_fail.py", line 32, in <module>
    complex_phase = phase_delay(lm, uvw, frequency)
  File "/home/sperkins/venv/afr/lib/python3.7/site-packages/jax/interpreters/mlir.py", line 171, in ir_constants
    raise TypeError("No constant handler for type: {}".format(type(val)))
TypeError: No constant handler for type: <class 'numpy.complex64'>
```


Discussion:
Thanks for the report. I can reproduce with jaxlib 0.1.76; the jax version doesn't appear to matter.
Shorter repro:
```python
import numpy as np
import jax
import jax.numpy as jnp


@jax.jit
def f():
  return jnp.exp(np.complex64(1j))
f()
```
I think the issue is that complex dtypes were left out of the list here: https://github.com/google/jax/blob/0382a6a04eddd7506a4ef6bb0c93f0f660ee3df6/jax/interpreters/mlir.py#L243-L247

Resolving patch:
diff --git a/jax/interpreters/mlir.py b/jax/interpreters/mlir.py
--- a/jax/interpreters/mlir.py
+++ b/jax/interpreters/mlir.py
@@ -241,9 +241,10 @@ def _ndarray_constant_handler(val: np.ndarray, canonicalize_types
 register_constant_handler(np.ndarray, _ndarray_constant_handler)
 
 for _scalar_type in [np.int8, np.int16, np.int32, np.int64,
-                    np.uint8, np.uint16, np.uint32, np.uint64,
-                    np.float16, np.float32, np.float64,
-                    np.bool_, np.longlong, dtypes.bfloat16]:
+                     np.uint8, np.uint16, np.uint32, np.uint64,
+                     np.float16, np.float32, np.float64,
+                     np.complex64, np.complex128,
+                     np.bool_, np.longlong, dtypes.bfloat16]:
   register_constant_handler(_scalar_type, _ndarray_constant_handler)
 
 def _python_scalar_handler(dtype, val, canonicalize_dtypes):


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -3120,6 +3120,16 @@ def f(_):
     expected = jnp.arange(1) + 1
     self.assertAllClose(ans, expected)
 
+  @parameterized.named_parameters([
+      {"testcase_name": f"{dtype.__name__}", "dtype": dtype}
+      for dtype in jtu.dtypes.all])
+  def test_constant_handlers(self, dtype):
+    # https://github.com/google/jax/issues/9380
+    @jax.jit
+    def f():
+      return jnp.exp(dtype(0))
+    f()  # doesn't error
+
   def test_large_python_ints(self):
     with self.assertRaises(OverflowError):
       jnp.multiply(2 ** 100, 3.)