From 165065756f178a1fc608877e96a180f12d3d3c3c Mon Sep 17 00:00:00 2001
From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com>
Date: Sat, 24 Sep 2022 18:27:22 -0700
Subject: [PATCH] Static type-checking fix (#34)

---
 jaxtyping/__init__.py | 20 ++++++++++++--------
 1 file changed, 12 insertions(+), 8 deletions(-)

diff --git a/jaxtyping/__init__.py b/jaxtyping/__init__.py
index ef39d1e..09e1a5b 100644
--- a/jaxtyping/__init__.py
+++ b/jaxtyping/__init__.py
@@ -20,14 +20,18 @@
 import typing
 
 
-if getattr(typing, "GENERATING_DOCUMENTATION", False):
-
-    class Array:
-        pass
-
-    Array.__module__ = "builtins"
-else:
+if typing.TYPE_CHECKING:
+    # type checkers don't know which branch below will be executed
     from jax.numpy import ndarray as Array
+else:
+    if getattr(typing, "GENERATING_DOCUMENTATION", False):
+
+        class Array:
+            pass
+
+        Array.__module__ = "builtins"
+    else:
+        from jax.numpy import ndarray as Array
 
 from .array_types import (
     AbstractArray,
@@ -63,4 +67,4 @@ from .import_hook import install_import_hook
 from .pytree_type import PyTree
 
 
-__version__ = "0.2.5"
+__version__ = "0.2.6"