[typing] Use ParamSpec in JIT annotation #14688

NeilGirdhar · 2023-02-26T02:15:02Z

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 · 2023-02-26T03:30:17Z

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 · 2023-02-26T03:37:56Z

setup.py

@@ -67,6 +67,7 @@ def generate_proto(source):
        'numpy>=1.20',
        'opt_einsum',
        'scipy>=1.5',
+        'typing_extensions>=4.5.0',


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

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

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.

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

@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 · 2023-02-26T12:36:58Z

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.

.pre-commit-config.yaml

jakevdp · 2023-02-26T13:47:58Z

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 · 2023-02-26T16:08:29Z

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

NeilGirdhar · 2023-02-26T19:09:22Z

@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 · 2023-08-16T21:07:42Z

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 · 2023-08-16T21:27:47Z

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 · 2023-08-16T21:58:18Z

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

NeilGirdhar · 2023-10-10T19:47:18Z

@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 · 2023-10-10T19:52:47Z

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 · 2023-10-11T19:55:40Z

#18066 updates mypy to v1.6.0

NeilGirdhar · 2023-10-11T21:12:46Z

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

NeilGirdhar · 2023-10-11T21:16:14Z

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.

jakevdp · 2023-11-03T22:51:59Z

jax/_src/numpy/lax_numpy.py

+@overload
+def atleast_1d(x: ArrayLike, y: ArrayLike, /, *arys: ArrayLike) -> list[Array]:
+  ...
+


Split out these changes in a separate PR?

Yes, okay, I will do that when we get close to having this checked in. This PR will expose a lot of typing errors as code that's jitted becomes typed.

Factored out as requested: #18395

NeilGirdhar · 2023-11-03T22:56:07Z

Seems like maybe a bad interaction with functools.partial?

That makes sense. I'll look at it as soon as I have more time.

NeilGirdhar · 2023-11-04T11:29:14Z

Filed a MyPy bug python/mypy#16404

NeilGirdhar · 2023-11-04T20:34:55Z

@jakevdp I think we should wait one more release of MyPy, which should fix the above errors. After that, although this change will not help MyPy users (the inferred type of a jit-decorated function will just be Callable), this change won't hurt, and it will help Pyright users. What do you think?

jakevdp · 2023-11-06T20:38:35Z

@jakevdp I think we should wait one more release of MyPy

Sounds good

NeilGirdhar · 2024-03-13T23:18:59Z

@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 · 2024-03-14T00:00:21Z

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 · 2024-03-14T00:09:09Z

@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 · 2024-03-14T00:32:24Z

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

jakevdp · 2024-03-14T12:50:07Z

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 · 2024-03-14T14:30:46Z

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

NeilGirdhar changed the title ~~Use ParamSpec in jit annotation; bump MyPy to 1.0~~ [typing] Use ParamSpec in jit annotation; bump MyPy to 1.0 Feb 26, 2023

NeilGirdhar force-pushed the jit_annotation branch from ba45b45 to 353bba1 Compare February 26, 2023 03:30

yashk2810 reviewed Feb 26, 2023

View reviewed changes

NeilGirdhar force-pushed the jit_annotation branch 2 times, most recently from c877b16 to ee60835 Compare February 26, 2023 11:50

NeilGirdhar marked this pull request as draft February 26, 2023 12:31

jakevdp reviewed Feb 26, 2023

View reviewed changes

.pre-commit-config.yaml Outdated Show resolved Hide resolved

hawkinsp assigned jakevdp Mar 10, 2023

NeilGirdhar force-pushed the jit_annotation branch 3 times, most recently from aaec678 to 0ffe6c8 Compare August 16, 2023 21:50

NeilGirdhar force-pushed the jit_annotation branch 2 times, most recently from ed78ba5 to 0e7d881 Compare August 18, 2023 00:41

jakevdp mentioned this pull request Aug 19, 2023

Functions decorated with @util._wraps(...) lose their type annotations. #16863

Closed

NeilGirdhar force-pushed the jit_annotation branch 2 times, most recently from ee38fd6 to 534cae8 Compare October 11, 2023 21:10

NeilGirdhar marked this pull request as ready for review October 11, 2023 21:11

jakevdp reviewed Nov 3, 2023

View reviewed changes

NeilGirdhar force-pushed the jit_annotation branch from 9816dad to 9f838ab Compare November 5, 2023 14:05

NeilGirdhar mentioned this pull request Nov 5, 2023

Fix annotations #18395

Merged

NeilGirdhar force-pushed the jit_annotation branch from 9f838ab to 91b4f1a Compare February 21, 2024 16:11

NeilGirdhar force-pushed the jit_annotation branch 2 times, most recently from 621d9c1 to bdc0d85 Compare March 13, 2024 23:15

jakevdp approved these changes Mar 13, 2024

View reviewed changes

google-ml-butler bot added kokoro:force-run pull ready Ready for copybara import and testing labels Mar 13, 2024

kokoro-team removed the kokoro:force-run label Mar 13, 2024

NeilGirdhar force-pushed the jit_annotation branch 3 times, most recently from c065282 to fcfd29c Compare March 16, 2024 09:10

NeilGirdhar force-pushed the jit_annotation branch from fcfd29c to b743a69 Compare March 25, 2024 10:50

NeilGirdhar added 2 commits March 25, 2024 06:51

Use ParamSpec in jit annotation

c0f0514

Make typing_extensions a dev-dependency

617850d

NeilGirdhar force-pushed the jit_annotation branch from b743a69 to 617850d Compare March 25, 2024 10:51

jakevdp mentioned this pull request May 8, 2024

typing: add annotations to util.cache & util.memoize #21131

Closed

This was referenced Sep 18, 2024

Improve typing of jax.jit #23720

Open

Typing improvement: Preserve the wrapped callable's signature in jax.jit #23719

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[typing] Use ParamSpec in JIT annotation #14688

[typing] Use ParamSpec in JIT annotation #14688

NeilGirdhar commented Feb 26, 2023 •

edited

Loading

NeilGirdhar commented Feb 26, 2023

yashk2810 Feb 26, 2023

NeilGirdhar Feb 26, 2023

NeilGirdhar Feb 26, 2023 •

edited

Loading

jakevdp Feb 26, 2023

NeilGirdhar Feb 26, 2023

NeilGirdhar commented Feb 26, 2023 •

edited

Loading

jakevdp commented Feb 26, 2023

JesseFarebro commented Feb 26, 2023

NeilGirdhar commented Feb 26, 2023 •

edited

Loading

NeilGirdhar commented Aug 16, 2023 •

edited

Loading

jakevdp commented Aug 16, 2023

NeilGirdhar commented Aug 16, 2023

NeilGirdhar commented Oct 10, 2023 •

edited

Loading

jakevdp commented Oct 10, 2023 •

edited

Loading

jakevdp commented Oct 11, 2023

NeilGirdhar commented Oct 11, 2023

NeilGirdhar commented Oct 11, 2023 •

edited

Loading

jakevdp Nov 3, 2023

NeilGirdhar Nov 3, 2023

NeilGirdhar Nov 5, 2023

NeilGirdhar commented Nov 3, 2023

NeilGirdhar commented Nov 4, 2023

NeilGirdhar commented Nov 4, 2023 •

edited

Loading

jakevdp commented Nov 6, 2023

NeilGirdhar commented Mar 13, 2024

jakevdp left a comment

jakevdp commented Mar 14, 2024

NeilGirdhar commented Mar 14, 2024 •

edited

Loading

jakevdp commented Mar 14, 2024

jakevdp commented Mar 14, 2024

NeilGirdhar commented Mar 14, 2024

[typing] Use ParamSpec in JIT annotation #14688

Are you sure you want to change the base?

[typing] Use ParamSpec in JIT annotation #14688

Conversation

NeilGirdhar commented Feb 26, 2023 • edited Loading

NeilGirdhar commented Feb 26, 2023

yashk2810 Feb 26, 2023

Choose a reason for hiding this comment

NeilGirdhar Feb 26, 2023

Choose a reason for hiding this comment

NeilGirdhar Feb 26, 2023 • edited Loading

Choose a reason for hiding this comment

jakevdp Feb 26, 2023

Choose a reason for hiding this comment

NeilGirdhar Feb 26, 2023

Choose a reason for hiding this comment

NeilGirdhar commented Feb 26, 2023 • edited Loading

jakevdp commented Feb 26, 2023

JesseFarebro commented Feb 26, 2023

NeilGirdhar commented Feb 26, 2023 • edited Loading

NeilGirdhar commented Aug 16, 2023 • edited Loading

jakevdp commented Aug 16, 2023

NeilGirdhar commented Aug 16, 2023

NeilGirdhar commented Oct 10, 2023 • edited Loading

jakevdp commented Oct 10, 2023 • edited Loading

jakevdp commented Oct 11, 2023

NeilGirdhar commented Oct 11, 2023

NeilGirdhar commented Oct 11, 2023 • edited Loading

jakevdp Nov 3, 2023

Choose a reason for hiding this comment

NeilGirdhar Nov 3, 2023

Choose a reason for hiding this comment

NeilGirdhar Nov 5, 2023

Choose a reason for hiding this comment

NeilGirdhar commented Nov 3, 2023

NeilGirdhar commented Nov 4, 2023

NeilGirdhar commented Nov 4, 2023 • edited Loading

jakevdp commented Nov 6, 2023

NeilGirdhar commented Mar 13, 2024

jakevdp left a comment

Choose a reason for hiding this comment

jakevdp commented Mar 14, 2024

NeilGirdhar commented Mar 14, 2024 • edited Loading

jakevdp commented Mar 14, 2024

jakevdp commented Mar 14, 2024

NeilGirdhar commented Mar 14, 2024

NeilGirdhar commented Feb 26, 2023 •

edited

Loading

NeilGirdhar Feb 26, 2023 •

edited

Loading

NeilGirdhar commented Feb 26, 2023 •

edited

Loading

NeilGirdhar commented Feb 26, 2023 •

edited

Loading

NeilGirdhar commented Aug 16, 2023 •

edited

Loading

NeilGirdhar commented Oct 10, 2023 •

edited

Loading

jakevdp commented Oct 10, 2023 •

edited

Loading

NeilGirdhar commented Oct 11, 2023 •

edited

Loading

NeilGirdhar commented Nov 4, 2023 •

edited

Loading

NeilGirdhar commented Mar 14, 2024 •

edited

Loading