From 6202dcc639fc1a12e229e24899a84145f243a1af Mon Sep 17 00:00:00 2001
From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com>
Date: Mon, 19 Sep 2022 23:30:29 -0700
Subject: [PATCH] doc fix (#26)

---
 jaxtyping/__init__.py    | 14 ++++++++++++--
 jaxtyping/array_types.py |  8 ++++++--
 2 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/jaxtyping/__init__.py b/jaxtyping/__init__.py
index c0b270f..eae9389 100644
--- a/jaxtyping/__init__.py
+++ b/jaxtyping/__init__.py
@@ -17,7 +17,17 @@
 # IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
-from jax.numpy import ndarray as Array
+import typing
+
+
+if getattr(typing, "GENERATING_DOCUMENTATION", False):
+
+    class Array:
+        pass
+
+    Array.__module__ = "builtins"
+else:
+    from jax.numpy import ndarray as Array
 
 from .array_types import (
     AbstractArray,
@@ -53,4 +63,4 @@ from .import_hook import install_import_hook
 from .pytree_type import PyTree
 
 
-__version__ = "0.2.2"
+__version__ = "0.2.3"
diff --git a/jaxtyping/array_types.py b/jaxtyping/array_types.py
index c03a09a..22dea71 100644
--- a/jaxtyping/array_types.py
+++ b/jaxtyping/array_types.py
@@ -19,6 +19,7 @@
 
 import enum
 import functools as ft
+import typing
 from typing import Any, Dict, List, NoReturn, Optional, Tuple, TYPE_CHECKING, Union
 from typing_extensions import Literal
 
@@ -385,7 +386,7 @@ class _MetaAbstractDtype(type):
         if _array_name_format == "dtype_and_shape":
             name = f"{cls.__name__}[{array_type.__name__}, '{dim_str}']"
         elif _array_name_format == "array":
-            name = "Array"
+            name = array_type.__name__
         else:
             raise ValueError(f"array_name_format {_array_name_format} not recognised")
         out = _MetaAbstractArray(
@@ -398,7 +399,10 @@ class _MetaAbstractDtype(type):
                 index_variadic=index_variadic,
             ),
         )
-        out.__module__ = "jaxtyping"
+        if getattr(typing, "GENERATING_DOCUMENTATION", False):
+            out.__module__ = "builtins"
+        else:
+            out.__module__ = "jaxtyping"
         return out