ScalarLike now dispatches to jaxtyping.ArrayLike rather than jax.typing.ArrayLike. This causes better behaviour during docgen.

Author: patrick-kidger

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