Use equality for extended dtypes in isclose and add test for `jax.r…

…andom.key`.

jax/_src/numpy/lax_numpy.py

@@ -938,6 +938,11 @@ def _moveaxis(a: Array, source: tuple[int, ...], destination: tuple[int, ...]) -
 @partial(jit, static_argnames=('equal_nan',))
 def isclose(a: ArrayLike, b: ArrayLike, rtol: ArrayLike = 1e-05, atol: ArrayLike = 1e-08,
             equal_nan: bool = False) -> Array:
+  a, b = util.promote_args("isclose", a, b)
+  dtype = _dtype(a)
+  if dtypes.issubdtype(dtype, dtypes.extended):
+    return lax.eq(a, b)
+
   a, b = util.promote_args_inexact("isclose", a, b)
   dtype = _dtype(a)
   if issubdtype(dtype, complexfloating):

tests/lax_numpy_test.py

@@ -3533,6 +3533,8 @@ def testIsClose(self):
         self.assertTrue(jnp.all(jnp.equal(result_np, result_jit)))
 
     self.assertEqual(np.isclose(6, 10, rtol=0.5), jnp.isclose(6, 10, rtol=0.5))
+    key = jax.random.key(0)
+    self.assertTrue(jnp.isclose(key, key))
 
   @jtu.sample_product(
     x=[1, [1], [1, 1 + 1E-4], [1, np.nan]],