From a6ab6c0d28a5209e08c31791a1a20dc6cace6697 Mon Sep 17 00:00:00 2001
From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com>
Date: Thu, 29 Jun 2023 18:39:22 -0700
Subject: [PATCH] Have PRNGKeyArray appear correctly in downstream
 documentation

---
 jaxtyping/__init__.py | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/jaxtyping/__init__.py b/jaxtyping/__init__.py
index 4f8e628..5714917 100644
--- a/jaxtyping/__init__.py
+++ b/jaxtyping/__init__.py
@@ -188,7 +188,17 @@ if typing.TYPE_CHECKING:
 
     from ._indirection import Scalar as Scalar, ScalarLike as ScalarLike
 elif has_jax:
-    from ._array_types import PRNGKeyArray, Scalar, ScalarLike  # noqa: F401
+    from ._array_types import Scalar, ScalarLike  # noqa: F401
+
+    if getattr(typing, "GENERATING_DOCUMENTATION", False):
+        # That is, we're generating some downstream documentation, not the jaxtyping
+        # documentation itself.
+        class PRNGKeyArray:
+            pass
+
+        PRNGKeyArray.__module__ = "builtins"
+    else:
+        from ._array_types import PRNGKeyArray
 
 del has_jax