From 58600d3fe0ef8780e1262b84240344c95a048108 Mon Sep 17 00:00:00 2001
From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com>
Date: Sat, 28 Oct 2023 13:51:15 -0700
Subject: [PATCH] Fixed new-style PRNG keys.

---
 jaxtyping/_array_types.py | 26 +++-----------------------
 test/test_array.py        | 30 ++++++++++++++++++++++++++++--
 2 files changed, 31 insertions(+), 25 deletions(-)

diff --git a/jaxtyping/_array_types.py b/jaxtyping/_array_types.py
index 9214014..edee1da 100644
--- a/jaxtyping/_array_types.py
+++ b/jaxtyping/_array_types.py
@@ -156,32 +156,12 @@ def _check_dims(
     return True
 
 
-def _is_jax_extended_dtype(dtype: Any) -> bool:
-    if not has_jax:
-        return False
-    try:
-        is_dtype = issubclass(dtype, jax.numpy.generic)
-    except TypeError:
-        # `dtype` not a class
-        return False
-    else:
-        if is_dtype:
-            if hasattr(jax.dtypes, "extended"):  # jax>=0.4.14
-                return jax.numpy.issubdtype(dtype, jax.dtypes.extended)
-            else:  # jax<=0.4.13
-                return jax.core.is_opaque_dtype(dtype)
-        else:
-            return False
-
-
 class _MetaAbstractArray(type):
     def __instancecheck__(cls, obj):
         if not isinstance(obj, cls.array_type):
             return False
 
-        if _is_jax_extended_dtype(obj.dtype):
-            dtype = str(obj.dtype)
-        elif hasattr(obj.dtype, "type") and hasattr(obj.dtype.type, "__name__"):
+        if hasattr(obj.dtype, "type") and hasattr(obj.dtype.type, "__name__"):
             # JAX, numpy
             dtype = obj.dtype.type.__name__
         elif hasattr(obj.dtype, "as_numpy_dtype"):
@@ -646,6 +626,7 @@ class AbstractDtype(metaclass=_MetaAbstractDtype):
         cls.dtypes = dtypes
 
 
+_prng_key = "prng_key"
 _bool = "bool"
 _bool_ = "bool_"
 _uint8 = "uint8"
@@ -714,8 +695,7 @@ Num = _make_dtype(uints + ints + floats + complexes, "Num")
 Shaped = _make_dtype(_any_dtype, "Shaped")
 
 if has_jax:
-    _key_regex = re.compile(r"^key<\w+>$")
-    Key = _make_dtype(_key_regex, "Key")
+    Key = _make_dtype(_prng_key, "Key")
     # New-style `jax.random.key` have scalar shape and dtype `key<foo>`.
     # Old-style `jax.random.PRNGKey` have shape `(2,)` and dtype `uint32`.
     PRNGKeyArray = Union[Key[jax.Array, ""], UInt32[jax.Array, "2"]]
diff --git a/test/test_array.py b/test/test_array.py
index 52b6e99..3d7c620 100644
--- a/test/test_array.py
+++ b/test/test_array.py
@@ -33,7 +33,9 @@ from jaxtyping import (
     Bool,
     Float,
     Float32,
+    Key,
     PRNGKeyArray,
+    Scalar,
     Shaped,
 )
 
@@ -586,8 +588,8 @@ def test_key(jaxtyp, typecheck):
     def f(x: PRNGKeyArray):
         pass
 
-    x = jr.PRNGKey(0)
-    f(x)
+    f(jr.key(0))
+    f(jr.PRNGKey(0))
 
     with pytest.raises(ParamError):
         f(object())
@@ -599,6 +601,30 @@ def test_key(jaxtyp, typecheck):
         f(jnp.array(3.0))
 
 
+def test_key_dtype(jaxtyp, typecheck):
+    @jaxtyp(typecheck)
+    def f1(x: Key[Array, ""]):
+        pass
+
+    @jaxtyp(typecheck)
+    def f2(x: Key[Scalar, ""]):
+        pass
+
+    for f in (f1, f2):
+        f(jr.key(0))
+
+        with pytest.raises(ParamError):
+            f(jr.PRNGKey(0))
+        with pytest.raises(ParamError):
+            f(object())
+        with pytest.raises(ParamError):
+            f(1)
+        with pytest.raises(ParamError):
+            f(jnp.array(3))
+        with pytest.raises(ParamError):
+            f(jnp.array(3.0))
+
+
 def test_extension(jaxtyp, typecheck, getkey):
     X = Shaped[Array, "a b"]
     Y = Shaped[X, "c d"]