experiment: add jax.typing module and use it in jax.lax

[jakevdp]

Part of #12049

Changes:

• created jax.typing module, following roadmap in JEP: Type Annotations #11859
• moved core.Shape definition into jax.typing
• moved jax.numpy.ndarray definition into jax.typing
• applied these new types throughout jax.lax, following the style recommended in JEP: Type Annotations #11859
• fixed annotations in other files that were broken by the new jax.lax types (what a tangled web we weave...)

Some observations:

• When implementing annotations in jax.lax, I used ArrayLike for inputs and NDArray for outputs. Originally I had NDArray aliased to the ndarray metaclass, but because jax.lax has several functions that explicitly return Tracer types, it seemed necessary to make NDArray = Union[ndarray, Tracer].
• This change had wide-ranging ripple effects: any function annotated with -> jnp.ndarray and returning the result of a lax operation caused a typecheck failure (because Tracer is inompatible), so it meant that in many places I had to change -> jnp.ndarray to -> NDarray, which caused even more ripples.
• An alternative would be to use -> jnp.ndarray by convention (with the understanding that it implies compatiblity with tracers) and use #type: ignore statements where tracers are explicitly returned.

Happy to hear any thoughts/opinions on this.

[patrick-kidger]

I think folks are definitely going to get confused between google/jax.typing and google/jaxtyping. Plus I'm very wary of yet more sharding within the JAX ecosystem.

Obviously the former has more restricted aims, but wants to stay within the static type system.

Can we see if there's a reasonable way to bring both projects together? For example switched on the typing.TYPE_CHECKING flag.

(EDIT: jakevdp and I are discussing this offline, and think we have most-of-a-plan.)

[jakevdp]

Can we see if there's a reasonable way to bring both projects together? For example switched on the typing.TYPE_CHECKING flag.

I'm not entirely sure what you have in mind... Do you mean that jax should depend on jaxtyping for static type checks? If so I don't think this is a good idea (I discussed this briefly in the roadmap draft). In my view, what jaxtyping includes is far more complicated than the level of type annotations we want in the core library, at least for now.

[jakevdp]

For context, this draft PR is about answering the question "can we use simple annotations that are slightly better than Array = Any without breaking mypy and pytype". My understanding of jaxtyping is that it's far beyond the scope of that question.

[jakevdp]

(closing in favor of #12300)