From 356f5b7f7b90aa1eeeaae771d1cb4224e5f0c6a6 Mon Sep 17 00:00:00 2001
From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com>
Date: Thu, 1 Jun 2023 11:06:02 -0700
Subject: [PATCH] Build fixes

---
 .github/workflows/release.yml |  2 +-
 docs/api/pytree.md            |  4 +++-
 jaxtyping/__init__.py         | 10 +++++++++-
 3 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
index 642f945..6c42a1e 100644
--- a/.github/workflows/release.yml
+++ b/.github/workflows/release.yml
@@ -33,7 +33,7 @@ jobs:
         with:
             python-version: "3.11"
             test-script: |
-                python -m pip install pytest beartype equinox jaxlib cloudpickle
+                python -m pip install -r ${{ github.workspace }}/test/requirements.txt
                 python -m pip install torch --extra-index-url https://download.pytorch.org/whl/cpu
                 cp -r ${{ github.workspace }}/test ./test
                 pytest
diff --git a/docs/api/pytree.md b/docs/api/pytree.md
index ff8bdf7..2933fa5 100644
--- a/docs/api/pytree.md
+++ b/docs/api/pytree.md
@@ -5,8 +5,10 @@
         members:
             false
 
-`jaxtyping.PyTreeDef` is an alias for `jax.tree_util.PyTreeDef`, which is the type of the return from `jax.tree_util.tree_structure(...)`.
+---
 
 :::jaxtyping.PyTreeDef
 
+---
+
 Note that `jaxtyping.{PyTree, PyTreeDef}` are only available if JAX has been installed.
diff --git a/jaxtyping/__init__.py b/jaxtyping/__init__.py
index 1dbd5e0..5bd1371 100644
--- a/jaxtyping/__init__.py
+++ b/jaxtyping/__init__.py
@@ -153,7 +153,15 @@ if typing.TYPE_CHECKING:
     # anything. (I believe this is sometimes called `Unknown`.) Thus, this odd-looking
     # annotation, which static type checkers aren't smart enough to resolve.
 elif has_jax:
-    from jax.tree_util import PyTreeDef as PyTreeDef
+    if hasattr(typing, "GENERATING_DOCUMENTATION"):
+
+        class PyTreeDef:
+            """Alias for `jax.tree_util.PyTreeDef`, which is the type of the return
+            from `jax.tree_util.tree_structure(...)`.
+            """
+
+    else:
+        from jax.tree_util import PyTreeDef as PyTreeDef
 
     from ._pytree_type import PyTree as PyTree  # noqa: F401