From 0c596ff3739b7658912afaef2d0f8d4ed06f48ce Mon Sep 17 00:00:00 2001
From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com>
Date: Wed, 10 May 2023 17:16:42 -0700
Subject: [PATCH] Fix for non-JAX installations.

---
 jaxtyping/__init__.py | 4 +++-
 pyproject.toml        | 2 +-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/jaxtyping/__init__.py b/jaxtyping/__init__.py
index fc7ffe7..5080be3 100644
--- a/jaxtyping/__init__.py
+++ b/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:
diff --git a/pyproject.toml b/pyproject.toml
index c5c55c6..b673759 100644
--- 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"