From 1e5229c20ec0ab19c397226d30bb460ed89ef42e Mon Sep 17 00:00:00 2001
From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com>
Date: Fri, 22 Sep 2023 10:58:55 -0700
Subject: [PATCH] Should be more robust to jax/numpy/tensorflow version changes

---
 .pre-commit-config.yaml   |  1 +
 jaxtyping/_array_types.py | 17 +++++++++++++----
 test/requirements.txt     |  3 ++-
 test/test_tf_dtype.py     | 13 +++++++++++++
 4 files changed, 29 insertions(+), 5 deletions(-)
 create mode 100644 test/test_tf_dtype.py

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 4a0b6d9..fdd00ea 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -26,3 +26,4 @@ repos:
     rev: 'v0.0.255'
     hooks:
       - id: ruff
+        args: ["--fix"]
diff --git a/jaxtyping/_array_types.py b/jaxtyping/_array_types.py
index fe4167e..c66f621 100644
--- a/jaxtyping/_array_types.py
+++ b/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):
diff --git a/test/requirements.txt b/test/requirements.txt
index 0df2227..a81b6f9 100644
--- a/test/requirements.txt
+++ b/test/requirements.txt
@@ -1,7 +1,8 @@
 beartype
 cloudpickle
 equinox
+IPython
 jaxlib
 pytest
+tensorflow
 typeguard<3
-IPython
diff --git a/test/test_tf_dtype.py b/test/test_tf_dtype.py
new file mode 100644
index 0000000..44d78e1
--- /dev/null
+++ b/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)