JAX is no longer a hard dependency (to support e.g. PyTorch) (#50)

jaxtyping/__init__.py

@@ -18,12 +18,21 @@
 # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
 import typing
+import typing_extensions
+
+
+try:
+    import jax
+except ImportError:
+    has_jax = False
+else:
+    has_jax = True
+    del jax
 
 
 if typing.TYPE_CHECKING:
-    # type checkers don't know which branch below will be executed
     from jax.numpy import ndarray as Array
-else:
+elif has_jax:
     if getattr(typing, "GENERATING_DOCUMENTATION", False):
 
         class Array:
@@ -64,7 +73,17 @@ from .array_types import (
 )
 from .decorator import jaxtyped
 from .import_hook import install_import_hook
-from .pytree_type import PyTree
 
 
+if typing.TYPE_CHECKING:
+    _T = typing.TypeVar("_T")
+
+    class PyTree(typing_extensions.Protocol[_T]):
+        pass
+
+elif has_jax:
+    from .pytree_type import PyTree
+
+del has_jax
+
 __version__ = "0.2.8"

jaxtyping/pytree_type.py

@@ -19,8 +19,7 @@
 
 import functools as ft
 import typing
-from typing import Generic, TYPE_CHECKING, TypeVar
-from typing_extensions import Protocol
+from typing import Generic, TypeVar
 
 import jax.tree_util as jtu
 import typeguard
@@ -87,18 +86,11 @@ class _MetaSubscriptPyTree(type):
         return all(map(is_leaftype, leaves))
 
 
-if TYPE_CHECKING:
-    # Work around pytype bug #1288
-    # pytype: skip-file
-    class PyTree(Protocol[_T]):
-        pass
-
-else:
-    PyTree = _MetaPyTree("PyTree", (), {})
-    if getattr(typing, "GENERATING_DOCUMENTATION", False):
-        PyTree.__module__ = "builtins"
-    else:
-        PyTree.__module__ = "jaxtyping"
 # Can't do `class PyTree(Generic[_T]): ...` because we need to override the
 # instancecheck for PyTree[foo], but subclassing
 # `type(Generic[int])`, i.e. `typing._GenericAlias` is disallowed.
+PyTree = _MetaPyTree("PyTree", (), {})
+if getattr(typing, "GENERATING_DOCUMENTATION", False):
+    PyTree.__module__ = "builtins"
+else:
+    PyTree.__module__ = "jaxtyping"

setup.py

@@ -71,7 +71,6 @@ python_requires = "~=3.7"
 # https://github.com/explosion/confection/blob/main/setup.cfg#L33 used in colab
 
 install_requires = [
-    "jax>=0.3.4",
     "numpy>=1.20.0",
     "typeguard>=2.13.3",
     "typing_extensions>=3.7.4.1",