Build fixes

.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

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.

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