JEP: Type Annotations

[jakevdp]

This Jax Enhancement Proposal (JEP) introduces a roadmap for type annotations in JAX, addressing a number of goals, non-goals, and design decisions that need to be made.

Part of #12049.

[NeilGirdhar]

I'm really happy to see this level of care about type annotations! This is a great design document. I love the division into three levels.

Just a few comments that are hopefully useful:

Flexible annotation for decorators has been a long-standing issue in the mypy package, which was only recently resolved by the introduction of ParamSpec in PEP 612, available starting in Python 3.10.
Because JAX follows NEP 29, it cannot rely on Python 3.10 features until sometime after mid-2024.

Do we have to wait for Python 3.10? ParamSpec is available as part of typing-extensions. Jax already uses typing-extensions in various places.

When functions are decorated by jax transformations like jit, vmap, grad, etc. JAX will strip all annotations. The reason for this is that without the mechanisms of PEP 612 there is no good way to do otherwise, and ParamSpec will not be available for use until Python 3.10.

This is an important point. I have been waiting for MyPy to fix their ParamSpec bugs before submitting a fix for the annotation stripping. The MyPy team is rolling out fixes very often, so I'm not even sure, but the fix may already work.

and functions returning a shape should always be Tuple[int]

Minor (common) error, this should be Tuple[int, ...].

[jakevdp]

Do we have to wait for Python 3.10? ParamSpec is available as part of typing-extensions. Jax already uses typing-extensions in various places.

The sentence after the one you quoted mentions this possibility 😁

I think that given how in-flux this is (recently added, not fully supported by mypy, etc.) we can leave this kind of annotation for a later data once things have stabilized. What do you think?

[jakevdp]

@NeilGirdhar – I edited the doc a bit to address some of those comments - let me know what you think

[jakevdp]

@8bitmp3 Thanks for the comments! I addressed them in the most recent commit.

[NeilGirdhar]

The sentence after the one you quoted mentions this possibility grin

My mistake!!

I think that given how in-flux this is (recently added, not fully supported by mypy, etc.) we can leave this kind of annotation for a later data once things have stabilized. What do you think?

Yes, of course!

[froystig]

Nice improvements since my last read!

[jakevdp]

Thanks - I'm squashing the commits and going to merge.

Thanks all for the discussion & comments on this! I think we have a nice framework to use going forward.

[patrick-kidger]

LGTM! I like where we ended up on this.

[NeilGirdhar]

With the last Jax, release it looks like we're using jax.Array now! I'm just curious, is there a plan to make Array generic the same way that numpy arrays are generic? That is, support something like:

from jax import Array
BooleanArray = Array[jnp.bool_]
IntegralArray = Array[jnp.integer[Any]]
RealArray = Array[jnp.floating[Any]]
ComplexArray = Array[Union[jnp.floating[Any], jnp.complexfloating[Any, Any]]]
KeyArray = Array[jax.random.PRNGDtype]

[jakevdp]

It would be nice, but it's not currently in scope: variadic type generics are still not fully supported by mypy and other type checkers. This is briefly addressed here: https://jax.readthedocs.io/en/latest/jep/12049-type-annotations.html#avoid-unstable-typing-mechanisms

[NeilGirdhar]

variadic type generics

I agree with avoiding unstable typing mechanisms, but, if I understand correctly, you don't need variadic type generics to have a generic dtype. I think you only need them for the shape.

[patrick-kidger]

Unfortunately, using variadic generics for shape annotations doesn't really work. See my previous discussions on the topic here and here. (Both these links cover a lot of the same points.)

TL;DR: with a lot of effort it might be possible to get something very limited. But that's such an unattractive value proposition that no-one has been interested in making that happen.

If you want dtype/shape annotations then for the forseeable future I think your best bet is jaxtyping.

[NeilGirdhar]

Unfortunately, using variadic generics for shape annotations doesn't really work.

Right, but I'm not suggesting using shape annotations. I'm only wondering whether we could have dtype annotations, which numpy already exposes and both Mypy and Pyright are able to check. Dtype annotations have already saved me many times, so it would be nice to keep them as I switch over to jax.Array.