Quick fix for docs not generating correctly since the last jaxtyping …

…release
patrick-kidger · Mar 6, 2024 · f708d17 · f708d17
1 parent 5e2518c
commit f708d17
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/jaxtyping/__init__.py b/jaxtyping/__init__.py
@@ -145,6 +145,7 @@ class Array:
                         pass
 
                     Array.__module__ = "builtins"
+                    Array.__qualname__ = "Array"
                     return Array
                 else:
                     import jax
@@ -157,6 +158,7 @@ class ArrayLike:
                         pass
 
                     ArrayLike.__module__ = "builtins"
+                    ArrayLike.__qualname__ = "ArrayLike"
                     return ArrayLike
                 else:
                     import jax.typing
@@ -169,6 +171,7 @@ class PRNGKeyArray:
                         pass
 
                     PRNGKeyArray.__module__ = "builtins"
+                    PRNGKeyArray.__qualname__ = "PRNGKeyArray"
                     return PRNGKeyArray
                 else:
                     # New-style `jax.random.key` have scalar shape and dtype `key<foo>`.
@@ -216,6 +219,7 @@ class PyTreeDef:
                     if typing.GENERATING_DOCUMENTATION:
                         # Equinox etc. docs get just `PyTreeDef`.
                         # jaxtyping docs get `jaxtyping.PyTreeDef`.
+                        PyTreeDef.__qualname__ = "PyTreeDef"
                         PyTreeDef.__module__ = "builtins"
                     return PyTreeDef
                 else:

diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "jaxtyping"
-version = "0.2.26"
+version = "0.2.27"
 description = "Type annotations and runtime checking for shape and dtype of JAX arrays, and PyTrees."
 readme = "README.md"
 requires-python ="~=3.9"