Skip to content

Commit

Permalink
clear up formatting of discussion
Browse files Browse the repository at this point in the history
  • Loading branch information
jakevdp committed Sep 9, 2022
1 parent a42bf48 commit 2da540e
Showing 1 changed file with 14 additions and 11 deletions.
25 changes: 14 additions & 11 deletions docs/jep/12049-type-annotations.md
Original file line number Diff line number Diff line change
Expand Up @@ -239,7 +239,7 @@ Again, there are a couple mechanisms that could be used for 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.

#### Following NumPy's lead
### Following NumPy's lead

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:
Expand All @@ -250,11 +250,11 @@ Given this, it may be reasonable to conform to NumPy's precedent and implement t

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.

#### Unification
### Unifying instance checks and annotation

Another approach would be to unify type checking and annotation via override mechanisms mentioned above.

##### Option 1: Partial unification
### Option 1: Partial unification
A partial unification might look like this:

- `jax.Array` is the actual type of on-device arrays.
Expand All @@ -264,7 +264,7 @@ A partial unification might look like this:
In this approach, `jax.numpy.ndarray` would become a simple alias `jax.typing.Array` for backward compatibility.


##### Option 2: Full unification via overrides
#### 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.
Expand All @@ -273,7 +273,7 @@ Alternatively, we could opt for full unification via overrides:

Here, `jax.numpy.ndarray` would become a simple alias `jax.Array` for backward compatibility.

##### Option 3: Full unification via class hierarchy
#### 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
Expand All @@ -283,7 +283,7 @@ Finally, we could opt for full unification via restructuring of the class hierar
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
#### 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
Expand All @@ -293,12 +293,15 @@ We could appease OOP pedants by instead making `Tracer` and `Array` derive from
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
#### Evaluation

There is no perfect option here, but from the point of view of the user, Options 2 and 3 arguably present the most clear user-facing API: `jax.Array` is all you need to know.
Option 3 (`Tracer` as a subclass of `Array`) is perhaps the purer approach, but on the other hand it would require `Tracer` objects to carry all the baggage of `Array` objects (data buffers, sharding, devices, etc.) For this reason, I (@jakevdp) believe that Option 2 presents the best path forward.
It offers the least confusing API for users 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__`; but I'd consider this a technical hurdle rather than a blocker.
Considering the overall strengths and weaknesses of each potential approach:

- From a user perspective, the unified approaches (options 2 and 3) are arguably best, because they remove the cognitive overhead involved in remembering which objects to use for instance checks or annotations: `jax.Array` is all you need to know
- Between Option 2 and Option 3, the purer (in an OOP sense) apporach is arguably Option 3 (`Tracer` as a subclass of `Array`), but in other senses it breaks the inheritance model, because it would require `Tracer` objects to carry all the baggage of `Array` objects (data buffers, sharding, devices, etc.)
- Option 2 is less pure in an OOP sense, but it aligns more closely with the spirit of how JAX is designed, with `Tracer` objects duck-typing as `Arrays`, and using mechanisms that Python provides to support this kind of duck typing. 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__`; but I'd consider this a technical hurdle rather than a blocker.

With this in mind, we conclude that Option 2 presents the best path forward.

### Implementation Plan

Expand Down

0 comments on commit 2da540e

Please sign in to comment.