From d2785baceda117d4f9a4ee3cb0173f1cb6fbed35 Mon Sep 17 00:00:00 2001
From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com>
Date: Thu, 14 Sep 2023 20:09:23 -0700
Subject: [PATCH] Better suppotr for incomplete JAX installations. Supersedes
 #105.

---
 jaxtyping/__init__.py     | 10 +---------
 jaxtyping/_array_types.py |  6 +++++-
 2 files changed, 6 insertions(+), 10 deletions(-)

diff --git a/jaxtyping/__init__.py b/jaxtyping/__init__.py
index 5714917..68f9260 100644
--- a/jaxtyping/__init__.py
+++ b/jaxtyping/__init__.py
@@ -20,20 +20,12 @@
 import importlib.metadata
 import typing
 
-
-try:
-    import jax
-except ImportError:
-    has_jax = False
-else:
-    has_jax = True
-    del jax
-
 # First import some things as normal
 from ._array_types import (
     AbstractArray as AbstractArray,
     AbstractDtype as AbstractDtype,
     get_array_name_format as get_array_name_format,
+    has_jax,
     set_array_name_format as set_array_name_format,
 )
 from ._decorator import jaxtyped as jaxtyped
diff --git a/jaxtyping/_array_types.py b/jaxtyping/_array_types.py
index b1f2013..fe4167e 100644
--- a/jaxtyping/_array_types.py
+++ b/jaxtyping/_array_types.py
@@ -32,7 +32,11 @@ from ._decorator import storage
 
 try:
     import jax
-except ImportError:
+except (ImportError, RuntimeError, AttributeError):
+    # We catch `RuntimeError` as JAX will throw this if it's present, but unable to run
+    # on the current machine. This fails with this error.
+    # We catch `AttributeError` as the above then leaves the module in a partially
+    # initialised state, which causes subsequent imports to fail with this error.
     has_jax = False
 else:
     has_jax = True