Should be more robust to jax/numpy/tensorflow version changes

.pre-commit-config.yaml

@@ -26,3 +26,4 @@ repos:
     rev: 'v0.0.255'
     hooks:
       - id: ruff
+        args: ["--fix"]

jaxtyping/_array_types.py

@@ -144,10 +144,19 @@ def _check_dims(
 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)
+    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):

test/requirements.txt

@@ -1,7 +1,8 @@
 beartype
 cloudpickle
 equinox
+IPython
 jaxlib
 pytest
+tensorflow
 typeguard<3
-IPython

test/test_tf_dtype.py

@@ -0,0 +1,13 @@
+# Tensorflow dependency kept in a separate file, so that we can optionally exclude it
+# more easily.
+import tensorflow as tf
+
+from jaxtyping import UInt
+
+
+def test_tf_dtype():
+    x = tf.constant(1, dtype=tf.uint8)
+    y = tf.constant(1, dtype=tf.float32)
+    hint = UInt[tf.Tensor, "..."]
+    assert isinstance(x, hint)
+    assert not isinstance(y, hint)