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)