simplify representation of ScalarLike when generating documentation

patrick-kidger · Feb 8, 2025 · 3e5f1fc · 3e5f1fc
1 parent 6d2021f
commit 3e5f1fc
Showing 1 changed file with 11 additions and 2 deletions.
diff --git a/jaxtyping/__init__.py b/jaxtyping/__init__.py
@@ -198,9 +198,18 @@ class PRNGKeyArray:
 
                 return Shaped[jax.Array, ""]
             elif item == "ScalarLike":
-                import jax.typing
+                if getattr(typing, "GENERATING_DOCUMENTATION", False):
+
+                    class ScalarLike:
+                        pass
 
-                return Shaped[jax.typing.ArrayLike, ""]
+                    ScalarLike.__module__ = "builtins"
+                    ScalarLike.__qualname__ = "ScalarLike"
+                    return ScalarLike
+                else:
+                    import jax.typing
+
+                    return jax.typing.ArrayLike
             elif item == "PyTree":
                 from ._pytree_type import PyTree