-
Notifications
You must be signed in to change notification settings - Fork 2.9k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
MAINT Clean up leftover Array = Any
aliases in jax/_src/**.py
#17760
Conversation
737e6a2
to
68f0ab1
Compare
If we're going to land this, I think we should fix rather than silence the errors as part of this PR. |
I fixed two out of three errors. The remaining one is tricky, because it overlaps dynamic shapes and ragged tensors and is not documented enough for me to figure out what the "right" (intended) types are. |
efbbe2d
to
8ca39a1
Compare
8ca39a1
to
86f6f07
Compare
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Looks good! Another couple small comments.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Last thing before merging – can you please squash your changes into a single commit?
e6ad41c
to
d5527df
Compare
ce86a3c
to
5c8586e
Compare
I had to revert to using `Any` for `RaggedAxis.ragged_axes` because pytype found more latent type errors, which require the understanding of ragedness and dynamic shapes internals to fix properly.
… in JAX Some JAX internal used Any instead of Array or in their type annotations. jax-ml/jax#17760 changed these to alias jax.Array and uncovered type errors fixed here. PiperOrigin-RevId: 569885017
… in JAX Some JAX internal used Any instead of Array or in their type annotations. jax-ml/jax#17760 changed these to alias jax.Array and uncovered type errors fixed here. PiperOrigin-RevId: 569885017
… in JAX Some JAX internal used Any instead of Array or in their type annotations. jax-ml/jax#17760 changed these to alias jax.Array and uncovered type errors fixed here. PiperOrigin-RevId: 569979327
… in JAX Some JAX internal used Any instead of Array or in their type annotations. jax-ml/jax#17760 changed these to alias jax.Array and uncovered type errors fixed here. PiperOrigin-RevId: 569979327
… in JAX Some JAX internal used Any instead of Array or in their type annotations. jax-ml/jax#17760 changed these to alias jax.Array and uncovered type errors fixed here. PiperOrigin-RevId: 569994709
… in JAX Some JAX internal used Any instead of Array or in their type annotations. jax-ml/jax#17760 changed these to alias jax.Array and uncovered type errors fixed here. PiperOrigin-RevId: 570088709
… in JAX Some JAX internal used Any instead of Array or in their type annotations. jax-ml/jax#17760 changed these to alias jax.Array and uncovered type errors fixed here. PiperOrigin-RevId: 570090425
… in JAX Some JAX internal used Any instead of Array or in their type annotations. jax-ml/jax#17760 changed these to alias jax.Array and uncovered type errors fixed here. PiperOrigin-RevId: 570090425
… in JAX Some JAX internal used Any instead of Array or in their type annotations. jax-ml/jax#17760 changed these to alias jax.Array and uncovered type errors fixed here. PiperOrigin-RevId: 570128353
This seemingly minor change uncovered a few inconsistencies in existing code. I silenced the new type errors via ``type: ignore` and added a TODO to investigate later.