Fix for non-JAX installations.

jaxtyping/__init__.py

@@ -115,7 +115,6 @@ else:
         Int32 as Int32,
         Int64 as Int64,
         Integer as Integer,
-        Key as Key,
         Num as Num,
         Shaped as Shaped,
         UInt as UInt,
@@ -125,6 +124,9 @@ else:
         UInt64 as UInt64,
     )
 
+    if has_jax:
+        from .array_types import Key as Key
+
 
 # Now import PyTree
 if typing.TYPE_CHECKING:

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "jaxtyping"
-version = "0.2.17"
+version = "0.2.18"
 description = "Type annotations and runtime checking for shape and dtype of JAX arrays, and PyTrees."
 readme = "README.md"
 requires-python ="~=3.8"