clarify implementation regarding jax.Array

Author: jakevdp

docs/jep/12049-type-annotations.md

@@ -219,33 +219,41 @@ Despite the familiarity of numpy's API choices, the `ndarray` / `NDArray` / `Arr
 
 Python itself is slowly moving to a world of unifying instance and annotation types; for example, with class-level `getitem` support in Python 3.9, classes like `typing.Dict`, `typing.List`, `typing.Tuple`, etc. are now [deprecated](https://peps.python.org/pep-0585/#implementation) in favor of their true-typed counterparts, `dict`, `list`, and `tuple`.
 
-With this in mind, JAX could instead choose to use its (eventual) `jax.Array` type directly for both annotation and instance checks. For handling tracers, tracers could be supported using a construct like the following:
+With this in mind, JAX could instead choose to use its (eventual) `jax.Array` type directly for both annotation and instance checks. For handling tracers within type annotations, we could use a construct like the following:
 ```python
 if TYPE_CHECKING:
   Array = Union[jax.Array, jax.Tracer]
 else:
-  Array = jax.Array
+  Array = jax._src.array.Array
 ```
-For instance checks, we could use the same metaclass override we currently do in the case of `np.ndarray`. And if we would like to support more granular shape/dtype-specific annotations in the future, we could do so via Python 3.9 class-level `__getitem__`, similar to how `list[int]` and `dict[str, int]` work for Python 3.9 built-in types.
+
+For handling instance checks, we could use the same metaclass override for `jax.Array` that we currently do in the case of `jnp.ndarray`. And if we would like to support more granular shape/dtype-specific annotations in the future, we could support constructions like `jax.Array[(3, 4), int]` via Python 3.9 class-level `__getitem__`, similar to how `list[int]` and `dict[str, int]` work for Python 3.9 built-in types.
 
 Another advantage of this unification route (using `jax.Array` as a generic array annotation) is that it fits well with the approach used by the `jaxtyping` library, which uses `jaxtyping.Array[...]` as its core annotation type.
 
+Given these advantages, it seems like the unified `jax.Array` is the better option compared to splitting annotation and instance logic between `jax.Array`, `jax.typing.NDArray`, and `jnp.ndarray`.
+
 
 ### Implementation Plan
 
 To move forward with type annotations, we will do the following:
 
 1. Iterate on this JEP doc until developers and stakeholders are bought-in.
 
-2. Create a `jax.typing` module (a thin wrapper around `jax._src.typing` as is our usual pattern) and put in it the first level of simple types mentioned above:
+2. Create a private `jax._src.typing` (not providing any public APIs for now) and put in it the first level of simple types mentioned above:
 
-  - `Array` / `ArrayLike` (For now we can use `Array = Union[ndarray, Tracer]`, and eventually plan to move to a simpler `Array = jax.Array` with enhancements discussed above)
+  - `Array`: because `jax.Array` is not yet out of experimental, we'll define `jax._src.typing.Array` with the type annotation features that will eventually move to `jax.Array` when it has landed.
+  - Move the `jnp.ndarray` definition into `typing.ndarray`; continue aliasing it at its current location.
+  - `ArrayLike`: a Union of types valid as inputs to normal `jax.numpy` functions
   - `Dtype` / `DtypeLike` (Check on capitalization of the `t`: what do other projects use?)
   - `Shape` / `NamedShape` / `ShapeLike`
-  - Move the `jnp.ndarray` definition into `typing.ndarray`; continue aliasing it at its current location.
 
 3. Once this is implemented, use these new typing definitions to comprehensively annotate functions within `jax.lax` according to the guidelines above.
 
 4. Continue adding additional annotations one module at a time, focusing on public API functions.
 
+5. Once `jax.Array` has fully landed, migrate the features of `jax._src.typing.Array` to this class, and let `jax.numpy.ndarray = Array`.
+
+6. When all is finalized, create a public `jax.typing` module that makes the above types available to users, along with documentation of annotation best practices for code using JAX.
+
 We will track this work in {jax-issue}`#12049`, which this JEP is named for.