[typing] Use ParamSpec in JIT annotation

[NeilGirdhar]

This pull request would be a huge improvement for Jax users who use type checkers like MyPy or Pyright, which now support ParamSpec.

Consider:

from jax import Array, jit

@jit
def f(x: Array, y: Array) -> Array:
    pass

reveal_type(f.__call__)

Previous to this PR, users who decorate a function with jit lose all annotations of the method:

Pyright: Type of "f.__call__" is "(*args: Unknown, **kwargs: Unknown) -> Unknown"
MyPy: Revealed type is "def (*args: Any, **kwargs: Any) -> Any"

After this PR, we get:

Pyright: Type of "f.__call__" is "(x: Array, y: Array) -> Array"
MyPy: Revealed type is "def (x: Array, y: Array) -> Array"

This unblocks a lot of type checking.

Part of #12049 cc: @jakevdp

[NeilGirdhar]

Not sure what's going on with the type checking, but perhaps it's not running MyPy 1.0.1?

Also, I'm still working on fixing this for applying jit to methods, which currently works, but does not type check correctly.

[yashk2810]

What's the reason for adding this as a compulsory dependency?

[NeilGirdhar]

I could make it a dev-dependency. Would that be better?

[NeilGirdhar]

I made it a dev-dependency in a subsequence change (within this PR) so that you can assess whether you like that better.

Normally, typing-extensions is seen as a light dependency since it only contains typing annotations, so it won't realistically break any runtime code.

[jakevdp]

I think it's fine to add as a dependency, but I'm also curious why it's necessary

[NeilGirdhar]

@jakevdp We need ParamSpec, which was added in Python 3.10 (Jax's minimum version per NEP 29 is 3.8. We could make it a dev-dependency (as I've now done) although I think it complicates the code.

[NeilGirdhar]

After testing this a lot with my code, I think MyPy is still not ready to check code with complex ParamSpec usage (although Pyright is).

If we check something like this in, it may trip up MyPy quite a bit. On the other hand, it has a the amazing effect of exposing plenty Jax type annotations that were previously hidden. Plenty of functions in jax.numpy are stripped of perfectly good annotations by @jit.

Also, there are some limitations in Python typing wrt method decorators. I've posted about this on python/typing. If people are using jit on methods, then they will get errors, and there doesn't seem to be a universal fix.

What I will do is add typed decorators to my tjax library. Please let me know if you'd like to get this in to Jax sooner or if we should wait on it.

[jakevdp]

Thanks for looking at this! ParamSpec seems like it could be interesting; that said it still looks to be relatively unstable and probably the remarks here still apply.

[JesseFarebro]

Hasn't this already been discussed here: #10311?

[NeilGirdhar]

@JesseFarebro My mistake, I didn't find that when I searched! MyPy has progressed since then, but maybe still not enough.

@jakevdp Yes, fair enough. ParamSpec seems to work for more cases than before, but it seems to be causing some errors still. I'll check back in a few months to see how support is coming.

[NeilGirdhar]

Great news: Thanks to the recently merged python/mypy#15837, it appears that the main MyPy error with this pull request may have been solved: python/mypy#12169. Also, many of the MyPy errors that may have affected its usage may have been solved: python/mypy#11846, python/mypy#12986, python/mypy#14802.

[jakevdp]

We've updated mypy to 1.4.1 in the meantime, and #17147 bumps it to 1.5.0 – can you sync your PR to the current main branch?

[NeilGirdhar]

@jakevdp Done. FYI the MyPy pull I linked is merged, but it's not in any released MyPy yet.

[NeilGirdhar]

@jakevdp MyPy 1.6 is out today, and it may now support using ParamSpec as in this pull request. (See the section on ParamSpec improvements.)

Is there a process or plan for upgrading Jax to use MyPy 1.6? Looks like you guys have been upgrading regularly.

[jakevdp]

As soon as the new mypy version is mirrored at https://github.com/pre-commit/mirrors-mypy, we can bump the version in the pre-commit configuration here: https://github.com/google/jax/blob/5ed692809da29ff4fd9a89faee9af6d4b41c7056/.pre-commit-config.yaml#L30

It looks like the mirror update will happen automatically in about 12.5 hours: https://github.com/pre-commit/mirrors-mypy/blob/08cbc46b6e135adec84911b20e98e5bc52032152/.github/workflows/main.yml#L6

[jakevdp]

#18066 updates mypy to v1.6.0

[NeilGirdhar]

Okay, I've rebased this to take advantage of the new MyPy version

[NeilGirdhar]

Just FYI if you want to get rid of the typing_extensions dependency, one option would be to follow SPEC 0, and drop Python 3.9 (since you already support Python 3.12). Python 3.10 has ParamSpec.

[NeilGirdhar]

@jakevdp I've rebased this now. Is there a way to get the tests to run? python/mypy#1484 is still unsolved, so there may only be modest gains in MyPy.

[jakevdp]

Looks good! I'm going to approve and then pull it in to run internal pytype tests

[jakevdp]

Pytype is failing with this message:

File ".../jax/_src/api.py", line 100, in <module>: argument "covariant" to TypeVar not supported yet [not-supported-yet]

Seems this is a known issue (google/pytype#1471). Is there any way to do this kind of improvement without using covariant TypeVars?

[NeilGirdhar]

@jakevdp

I'm going to approve and then pull it in to run internal pytype tests'

Thanks for taking the time!

Is there any way to do this kind of improvement without using covariant TypeVars?

The return type does need to be covariant (MyPy gives an error without it). One option is to do this in two steps. Just keep the parameter specification annotation and remove the return type annotation.

It's funny because in Python 3.12, we won't need these markers at all as the type checker can infer them, although I think we'd have to use the new PEP 695 syntax.

Should I remove the annotations on the return type?

[jakevdp]

I'm honestly not sure what the best fix is here.

[jakevdp]

I think maybe the discussion in https://jax.readthedocs.io/en/latest/jep/12049-type-annotations.html#avoid-unstable-typing-mechanisms is still applicable.

[NeilGirdhar]

@jakevdp No worries. I thought we were just waiting for MyPy. I didn't realize that we were waiting for Pytype as well!

[NeilGirdhar]

@jakevdp Given PyType's announcement, looks like PyType will never support ParamSpec. Should we take another look at this?

[jakevdp]

Possibly, but maybe sometime later this fall. Our type checking story will likely be changing in the coming months (e.g. folks on the team are excited about ty) but it's too early to know exactly where we'll land.

[NeilGirdhar]

Sounds great. I'm also excited about ty!