Issue 1:
Please:

Check for duplicate issues.

Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:

The following test case 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.

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)

#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'>

Issue 2:
Please:

Check for duplicate issues.

Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:

Install jax v0.3.5 and CUDA-enabled jaxlib v0.3.0.

Run the test suite.

Observe the cuSparseTest.test_coo_sorted_indices_gpu_lowerings test to fail:

=================================== FAILURES ===================================
______________ cuSparseTest.test_coo_sorted_indices_gpu_lowerings ______________
[gw0] linux -- Python 3.9.11 /nix/store/bvzpppawf2naqv3qhxqv66jxaq3iaxyj-python3-3.9.11/bin/python3.9

self = <sparse_test.cuSparseTest testMethod=test_coo_sorted_indices_gpu_lowerings>

@unittest.skipIf(not GPU_LOWERING_ENABLED, "test requires cusparse/hipsparse")
@jtu.skip_on_devices("rocm") # TODO(rocm): see SWDEV-328107
def test_coo_sorted_indices_gpu_lowerings(self):
dtype = jnp.float32

mat = jnp.arange(12, dtype=dtype).reshape(4, 3)

mat_rows_sorted = sparse.COO.fromdense(mat)
self.assertTrue(mat_rows_sorted._rows_sorted)
self.assertFalse(mat_rows_sorted._cols_sorted)

mat_cols_sorted = sparse.COO.fromdense(mat.T).T
self.assertFalse(mat_cols_sorted._rows_sorted)
self.assertTrue(mat_cols_sorted._cols_sorted)

mat_unsorted = sparse.COO(mat_rows_sorted._bufs, shape=mat_rows_sorted.shape)
self.assertFalse(mat_unsorted._rows_sorted)
self.assertFalse(mat_unsorted._cols_sorted)

self.assertArraysEqual(mat, mat_rows_sorted._sort_rows().todense())
self.assertArraysEqual(mat, mat_cols_sorted._sort_rows().todense())
self.assertArraysEqual(mat, mat_unsorted._sort_rows().todense())

todense = jit(sparse.coo_todense)
with self.assertNoWarnings():
dense_rows_sorted = todense(mat_rows_sorted)
dense_cols_sorted = todense(mat_cols_sorted)
dense_unsorted = todense(mat_unsorted._sort_rows())
with self.assertWarnsRegex(sparse.CuSparseEfficiencyWarning, "coo_todense GPU lowering requires matrices with sorted rows.*"):
> dense_unsorted_fallback = todense(mat_unsorted)
E AssertionError: CuSparseEfficiencyWarning not triggered

tests/sparse_test.py:473: AssertionError
=========================== short test summary info ============================
FAILED tests/sparse_test.py::cuSparseTest::test_coo_sorted_indices_gpu_lowerings
========== 1 failed, 14145 passed, 1487 skipped in 165.49s (0:02:45) ===========
error: builder for '/nix/store/bk64c49429aj73wkz4dijlf15h0l2v9r-python3.9-jax-0.3.5.drv' failed with exit code 1;
last 10 log lines:
> dense_cols_sorted = todense(mat_cols_sorted)
> dense_unsorted = todense(mat_unsorted._sort_rows())
> with self.assertWarnsRegex(sparse.CuSparseEfficiencyWarning, "coo_todense GPU lowering requires matrices with sorted rows.*"):
> > dense_unsorted_fallback = todense(mat_unsorted)
> E AssertionError: CuSparseEfficiencyWarning not triggered
>
> tests/sparse_test.py:473: AssertionError
> =========================== short test summary info ============================
> FAILED tests/sparse_test.py::cuSparseTest::test_coo_sorted_indices_gpu_lowerings
> ========== 1 failed, 14145 passed, 1487 skipped in 165.49s (0:02:45) ===========

If applicable, include full error messages/tracebacks.

Why isn't this test skipped when running on a machine without access to a GPU?