revamp unification discussion

Author: jakevdp

docs/jep/12049-type-annotations.md

@@ -175,7 +175,7 @@ Note that these will in general be simpler than the equivalent protocols used in
 
 Conversely, outputs of functions and methods should be typed as strictly as possible: for example, for a JAX function that returns an array, the output should be annotated with `jnp.ndarray` rather than `ArrayLike`. Functions returning a dtype should always be annotated `np.dtype`, and functions returning a shape should always be `Tuple[int]` or a strictly-typed NamedShape equivalent. For this purpose, we will implement in {mod}`jax.typing` several strictly-typed analogs of the permissive types mentioned above, namely:
 
-- `NDArray` is effectively equivalent to `Union[Tracer, jnp.ndarray]` and should be used to annotate array outputs.
+- `Array` or `NDArray` (see below) for type annotation purposes is effectively equivalent to `Union[Tracer, jnp.ndarray]` and should be used to annotate array outputs.
 - `DType` is an alias of `np.dtype`, perhaps with the ability to also represent key types and other generalizations used within JAX.
 - `Shape` is essentially `Tuple[int, ...]`, perhaps with some additional flexibilty to account for dynamic shapes.
 - `NamedShape` is an extension of `Shape` that allows for named shapes as used internall in JAX.
@@ -197,40 +197,104 @@ We will revisit this question in the future once support for such features stabi
 
 Similarly, regarding [jaxtyping](http://github.com/google/jaxtyping) project, despite its promise as a project, at this point we consider its goals beyond the scope of the simpler level of type annotations we would like for the jax core project. This is a decision we could revisit at a future date.
 
-### Design Considerations
+### `Array`/`NDArray` Design Considerations
 
-#### Following NumPy's Lead
+Type annotation of arrays in JAX poses a unique challenge because of JAX's extensive use of duck-typing, i.e. passing and returning `Tracer` objects in place actual arrays within jax transformations.
+This becomes increasingly confusing because objects used for type annotation often overlap with objects used for runtime instance checking, and may or may not correspond to the actual type hierarchy of the objects in question.
+For JAX, we need to provide duck-typed objects for use in two contexts: **static type annotations** and **runtime instance checks**.
 
-One possible model for JAX's type annotations is those provided by the [`numpy.typing`](https://numpy.org/devdocs/reference/typing.html) module.
-An advantage of this approach is that it is that the API is already familiar to many; most notably `numpy.typing` provides:
+The following discussion will assume that `jax.Array` is the type used to represent on-device arrays, which is not yet the case but will be once the work in {jax-issue}`#12016` is complete.
 
-- `NDArray`, which is a generic array type used only for type annotations
-- `np.ndarray`, which is the array type used for instance checks, and also works for type annotations. In Python 3.9, can use constructs of the form `np.ndarray[shape, dtype]`, but this is currently poorly documented and it's unclear how well supported that is at the moment.
+1. **Static type annotations.**
+  We need to provide an object that can be used for duck-typed type annotations.
+  Assuming for the moment that we call this object `ArrayAnnotation`, we need a solution which satisfies `mypy` and `pytype` for a case like the following:
+  ```python
+  @jit
+  def f(x: ArrayAnnotation) -> ArrayAnnotation:
+    assert isinstance(x, core.Tracer)
+    return x
+  ```
+  This could be accomplished via a number of approaches, for example:
+  - Use a type union: `ArrayAnnotation = Union[Array, Tracer]`
+  - Create an interface file that declares `Tracer` and `Array` should be treated as subclasses of `ArrayAnnotation`.
+  - Restructure `Array` and `Tracer` so that `ArrayAnnotation` is a true base class of both.
+2. **Runtime instance checks.**
+  We also must provide an object that can be used for duck-typed runtime `isinstance` checks.
+  Assuming for the moment that we call this object `ArrayInstance`, we need a solution that passes the following runtime check:
+  ```python
+  def f(x):
+    return isinstance(x, ArrayInstance)
+  x = jnp.array([1, 2, 3])
+  assert f(x)       # x will be an array
+  assert jit(f)(x)  # x will be a tracer
+  ```
+  Again, there are a couple mechanisms that could be used for this:
+  - override `type(ArrayInstance).__instancecheck__` to return `True` for both `Array` and `Tracer` objects; this is how `jnp.ndarray` is currently implemented ([source](https://github.com/google/jax/blob/jax-v0.3.17/jax/_src/numpy/ndarray.py#L24-L49)).
+  - define `ArrayInstance` as an abstract base class and dynamically register it to `Array` and `Tracer`
+  - restructure `Array` and `Tracer` so that `ArrayInstance` is a true base class of both `Array` and `Tracer`
 
-JAX currently implements {class}`jax.numpy.ndarray` for use with runtime `isinstance` checks. It uses a metaclass override to ensure that tracers also will pass `isinstance(tracer, jnp.ndarray)`, but it does not currently have any mechanism to ensure that tracers will be valid for `jnp.ndarray` annotations. If we were to follow numpy's approach to type annotation, changes might look something like this:
+A decision we need to make is whether `ArrayAnnotation` and `ArrayInstance` should be the same or different objects. There is some precedent here; for example in the core Python language spec, `typing.Dict` and `typing.List` exist for the sake of annotation, while the built-in `dict` and `list` serve the purposes of instance checks.
+However, `Dict` and `List` are [deprecated](https://peps.python.org/pep-0585/#implementation) in newer Python versions in favor of using `dict` and `list` for both annotation and instance checks.
 
-- add `jax.typing.NDArray` for use with type annotations
-- add `TYPE_CHECKING` logic to ensure that the `jnp.ndarray` object can also be used as an annotation for tracers; and possibly add class-level `__getitem__` to match numpy's shape and dtype specification.
+#### Following NumPy's lead
 
-A potential point of confusion with this is that JAX arrays are not actually of type `jnp.ndarray`, but rather `DeviceArray` or `ShardedDeviceArray` (soon to be unified under a single `jax.Array` type; see {jax-issue}`#12016`).
-Following numpy's lead could result in confusion: we'd have `jax.Array`, `jax.numpy.ndarray`, and `jax.typing.NDArray`, each of which is useful in a particular subset of cases but not others (type identity, isinstance checks, type annotations, and tracer-compatible versions of the above).
-Despite the familiarity of numpy's API choices, the `ndarray` / `NDArray` / `Array` trichotomy may cause too much confusion
+In NumPy's case, `np.typing.NDArray` serves the purpose of type annotations, while `np.ndarray` serves the purpose of instance checks (as well as array type identity).
+Given this, it may be reasonable to conform to NumPy's precedent and implement the following:
 
-#### Choosing our own path: Unification
+- `jax.Array` is the actual type of on-device arrays.
+- `jax.typing.NDArray` is the object used for duck-typed array annotations.
+- `jax.numpy.ndarray` is the object used for duck-typed array instance checks.
 
-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`.
+This might feel somewhat natural to NumPy power-users, however this trifurcation would likely be a source of confusion: the choice of which to use for instance checks and annotations is not immediately clear.
 
-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._src.array.Array, jax.Tracer]
-else:
-  Array = jax._src.array.Array
-```
+#### Unification
+
+Another approach would be to unify type checking and annotation via override mechanisms mentioned above.
+
+##### Option 1: Partial unification
+A partial unification might look like this:
+
+- `jax.Array` is the actual type of on-device arrays.
+- `jax.typing.Array` is the object used for duck-typed array annotations (via `.pyi` interfaces on `Array` and `Tracer`).
+- `jax.typing.Array` is also the object used duck-typed instance checks (via an `__isinstance__` override in its metaclass)
+
+In this approach, `jax.numpy.ndarray` would become a simple alias `jax.typing.Array` for backward compatibility.
+
+
+##### Option 2: Full unification via overrides
+Alternatively, we could opt for full unification via overrides:
+
+- `jax.Array` is the actual type of on-device arrays.
+- `jax.Array` is also the object used for duck-typed array annotations (via a `.pyi` interface on `Tracer`)
+- `jax.Array` is also the object used for duck-typed instance checks (via an `__isinstance__` override in its metaclass)
 
-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, this would set us up to follow the conventions being developed in the [`jaxtyping`](https://github.com/google/jaxtyping/) project. Because `jaxtyping` uses `jaxtyping.Array[...]` as its core annotation type, unifying under `jax.Array` makes any potential future integration with that project more natural.
+Here, `jax.numpy.ndarray` would become a simple alias `jax.Array` for backward compatibility.
 
-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`.
+##### Option 3: Full unification via class hierarchy
+Finally, we could opt for full unification via restructuring of the class hierarchy and replacing duck-typing with OOP object hierarchies:
+
+- `jax.Array` is the actual type of on-device arrays
+- `jax.Array` is also the object used for array annotations, by ensuring that `Tracer` inherits from `jax.Array`
+- `jax.Array` is also the object used for instance checks, via the same mechanism
+
+Here `jnp.ndarray` could be an alias for `jax.Array`.
+This final approach is in some senses the most pure, but it may be challenging from an OOP design standpoint (`Tracer` *is a* `Array`?).
+
+##### Option 4: Parial unification via class hierarchy
+We could appease OOP pedants by instead making `Tracer` and `Array` derive from a common `ArrayBase` base class:
+
+- `jax.Array` is the actual type of on-device arrays
+- `ArrayBase` is the object used for array annotations
+- `ArrayBase` is also the object used for instance checks
+
+Here `jnp.ndarray` would be an alias for `ArrayBase`.
+This may be purer from an OOP perspective, but it reintroduces a bifurcation and the distinction between `Array` and `ArrayBase` for annotation and instance checks may become confusing.
+
+##### Evaluation
+
+There is no perfect option here, but weighing the pros and cons of these solutions, I (@jakevdp) believe that Option 4 presents the best path forward.
+It offers the least confusing API for users (`jax.Array` is the only object you need to worry about), and does not require any significant restructuring of our existing codepaths.
+There is one minor technical hurdle involved; that is that `jax.Array` will be defined in C++ via pybind11, and pybind11 currently [does not support](https://github.com/pybind/pybind11/issues/2696) custom metaclasses required for overriding `__instancecheck__`; nevertheless we should be able to work around this.
 
 ### Implementation Plan
 
@@ -240,17 +304,18 @@ To move forward with type annotations, we will do the following:
 
 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`: 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.
+  - `Array`: because `jax.Array` is not yet out of experimental, we'll define `jax._src.typing.Array` with the type annotation and instance-checking features that will eventually move to `jax.Array` when it has fully landed.
   - `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`
 
-3. Once this is implemented, use these new typing definitions to comprehensively annotate functions within `jax.lax` according to the guidelines above.
+3. When this is implemented, remove the existing definition of `jnp.ndarray` and set is as an alias of `jax._src.typing.Array`.
+
+4. As a test, 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. 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. 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.