From 13e6870fb831c39ec20aade620c131d74ecbfc5e Mon Sep 17 00:00:00 2001
From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com>
Date: Wed, 26 Jul 2023 09:30:40 -0700
Subject: [PATCH] Compatibility with JAX changes to opaque dtypes

---
 jaxtyping/_array_types.py | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/jaxtyping/_array_types.py b/jaxtyping/_array_types.py
index 8d4120d..8113829 100644
--- a/jaxtyping/_array_types.py
+++ b/jaxtyping/_array_types.py
@@ -137,12 +137,21 @@ def _check_dims(
     return True
 
 
+def _is_jax_extended_dtype(dtype: Any) -> bool:
+    if not has_jax:
+        return False
+    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)
+
+
 class _MetaAbstractArray(type):
     def __instancecheck__(cls, obj):
         if not isinstance(obj, cls.array_type):
             return False
 
-        if has_jax and jax.core.is_opaque_dtype(obj.dtype):
+        if _is_jax_extended_dtype(obj.dtype):
             dtype = str(obj.dtype)
         elif hasattr(obj.dtype, "type") and hasattr(obj.dtype.type, "__name__"):
             # JAX, numpy