Fix for non-JAX installations.

patrick-kidger · May 11, 2023 · 0c596ff · 0c596ff
1 parent 7934d2a
commit 0c596ff
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/jaxtyping/__init__.py b/jaxtyping/__init__.py
@@ -115,7 +115,6 @@ class ArrayLike:
         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 @@ class ArrayLike:
         UInt64 as UInt64,
     )
 
+    if has_jax:
+        from .array_types import Key as Key
+
 
 # Now import PyTree
 if typing.TYPE_CHECKING:

diff --git a/pyproject.toml b/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"