Repository: google/jax

Issue:
Type checking behaviour for RBG PRNG different when upgrading to jax v0.4.16
### Description

When upgrading to jax v0.4.16, the following `DeprecationWarning` is shown:

```
DeprecationWarning: jax.random.PRNGKeyArray is deprecated. Use jax.Array for annotations, and jax.dtypes.issubdtype(arr, jax.dtypes.prng_key) for runtime detection of typed prng keys.
```

This implies that the behaviour of `isinstance(key, jax.random.PRNGKeyArray)` and `jax.dtypes.issubdtype(key, jax.dtypes.prng_key)` should be the same. However, they are not the same when using RBG PRNG.

Code to reproduce:

```python
import os; os.environ['JAX_PLATFORMS'] = 'cpu'

import jax
import jax.random as rand

jax.config.update('jax_enable_custom_prng', True)
jax.config.update('jax_default_prng_impl', 'rbg')

key = rand.PRNGKey(3407)
print(isinstance(key, jax.random.PRNGKeyArray))
print(jax.dtypes.issubdtype(key, jax.dtypes.prng_key))
```

Output:

```
True
False
```


### What jax/jaxlib version are you using?

jax 0.4.16

### Which accelerator(s) are you using?

CPU

### Additional system info

Python 3.11.5, Ubuntu 20.04.6 LTS x86_64

### NVIDIA GPU info

_No response_


Discussion:
I think you meant this:
```python
jax.dtypes.issubdtype(key.dtype, jax.dtypes.prng_key)
```
With that change, the code should work as expected.
Looks like the snippet in the warning is wrong – I'll correct that.

You can read more about this in https://jax.readthedocs.io/en/latest/jep/9263-typed-keys.html
Fixed in #17709 – thanks for the report!

Resolving patch:
diff --git a/jax/_src/dtypes.py b/jax/_src/dtypes.py
--- a/jax/_src/dtypes.py
+++ b/jax/_src/dtypes.py
@@ -327,13 +327,15 @@ def issubdtype(a: DTypeLike, b: DTypeLike) -> bool:
   This is like :func:`numpy.issubdtype`, but can handle dtype extensions such as
   :obj:`jax.dtypes.bfloat16`.
   """
+  # Handle extended types & canonicalizes all concrete types to np.dtype instances.
   if isinstance(a, ExtendedDType):
     return _issubclass(a.type, b)
-  elif _issubclass(b, extended):
-    return False
-  # Canonicalizes all concrete types to np.dtype instances
   a = a if _is_typeclass(a) else np.dtype(a)
+
+  if _issubclass(b, extended):
+    return False
   b = b if _is_typeclass(b) else np.dtype(b)
+
   if isinstance(a, np.dtype):
     if a in _custom_float_dtypes:
       # Avoid implicitly casting list elements below to a dtype.
@@ -584,6 +586,8 @@ def dtype(x: Any, *, canonicalize: bool = False) -> DType:
     dt = python_scalar_dtypes[x]
   elif type(x) in python_scalar_dtypes:
     dt = python_scalar_dtypes[type(x)]
+  elif _issubclass(x, np.generic):
+    return np.dtype(x)
   elif issubdtype(getattr(x, 'dtype', None), extended):
     dt = x.dtype
   else:
diff --git a/jax/random.py b/jax/random.py
--- a/jax/random.py
+++ b/jax/random.py
@@ -194,12 +194,12 @@
     # Added September 13, 2023:
     "PRNGKeyArray": (
         "jax.random.PRNGKeyArray is deprecated. Use jax.Array for annotations, and "
-        "jax.dtypes.issubdtype(arr, jax.dtypes.prng_key) for runtime detection of "
+        "jax.dtypes.issubdtype(arr.dtype, jax.dtypes.prng_key) for runtime detection of "
         "typed prng keys.", _PRNGKeyArray
     ),
     "KeyArray": (
         "jax.random.KeyArray is deprecated. Use jax.Array for annotations, and "
-        "jax.dtypes.issubdtype(arr, jax.dtypes.prng_key) for runtime detection of "
+        "jax.dtypes.issubdtype(arr.dtype, jax.dtypes.prng_key) for runtime detection of "
         "typed prng keys.", _PRNGKeyArray
     ),
     # Added September 21, 2023


Test patch:
diff --git a/tests/dtypes_test.py b/tests/dtypes_test.py
--- a/tests/dtypes_test.py
+++ b/tests/dtypes_test.py
@@ -638,6 +638,7 @@ def testArrayRepr(self, dtype, weak_type):
       numpy_dtype_promotion=['strict', 'standard']
   )
   def testSafeToCast(self, input_dtype, output_dtype, numpy_dtype_promotion):
+    print(input_dtype, output_dtype)
     with jax.numpy_dtype_promotion(numpy_dtype_promotion):
       # First the special cases which are always safe:
       always_safe = (
diff --git a/tests/random_test.py b/tests/random_test.py
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -62,7 +62,7 @@ def _prng_key_as_array(key):
 def _maybe_unwrap(key):
   # TODO(frostig): remove some day when we deprecate "raw" key arrays
   unwrap = prng_internal.random_unwrap
-  return unwrap(key) if jnp.issubdtype(key, dtypes.prng_key) else key
+  return unwrap(key) if jnp.issubdtype(key.dtype, dtypes.prng_key) else key
 
 
 PRNG_IMPLS = [('threefry2x32', prng_internal.threefry_prng_impl),
@@ -1742,6 +1742,8 @@ def test_issubdtype(self):
     key = random.key(42)
     self.assertTrue(jnp.issubdtype(key.dtype, dtypes.prng_key))
     self.assertFalse(jnp.issubdtype(key.dtype, np.integer))
+    with self.assertRaisesRegex(TypeError, "Cannot interpret"):
+      jnp.issubdtype(key, dtypes.prng_key)
 
   @skipIf(not config.jax_enable_custom_prng, 'relies on typed key upgrade flag')
   def test_construction_upgrade_flag(self):