ScalarLike now dispatches to jaxtyping.ArrayLike rather than jax.typi…

…ng.ArrayLike. This causes better behaviour during docgen.
patrick-kidger · Feb 15, 2025 · d3c6b0f · d3c6b0f
1 parent 6813074
commit d3c6b0f
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/jaxtyping/__init__.py b/jaxtyping/__init__.py
@@ -198,9 +198,9 @@ class PRNGKeyArray:
 
                 return Shaped[jax.Array, ""]
             elif item == "ScalarLike":
-                import jax.typing
+                from . import ArrayLike
 
-                return Shaped[jax.typing.ArrayLike, ""]
+                return Shaped[ArrayLike, ""]
             elif item == "PyTree":
                 from ._pytree_type import PyTree