Skip to content
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

Merged
merged 1 commit into from
Oct 3, 2023

Conversation

superbobry
Copy link
Collaborator

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.

@jakevdp
Copy link
Collaborator

jakevdp commented Sep 25, 2023

If we're going to land this, I think we should fix rather than silence the errors as part of this PR.

@superbobry
Copy link
Collaborator Author

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.

jax/_src/interpreters/batching.py Outdated Show resolved Hide resolved
jax/_src/lax/convolution.py Outdated Show resolved Hide resolved
jax/_src/nn/functions.py Show resolved Hide resolved
jax/_src/nn/functions.py Outdated Show resolved Hide resolved
jax/_src/nn/functions.py Show resolved Hide resolved
Copy link
Collaborator

@jakevdp jakevdp left a 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.

jax/_src/nn/functions.py Outdated Show resolved Hide resolved
jax/_src/nn/functions.py Show resolved Hide resolved
@google-ml-butler google-ml-butler bot added kokoro:force-run pull ready Ready for copybara import and testing labels Sep 27, 2023
Copy link
Collaborator

@jakevdp jakevdp left a 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?

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.
copybara-service bot pushed a commit to google-deepmind/clrs that referenced this pull request Oct 1, 2023
… 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
copybara-service bot pushed a commit to google-deepmind/clrs that referenced this pull request Oct 2, 2023
… 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
copybara-service bot pushed a commit to google-deepmind/clrs that referenced this pull request Oct 2, 2023
… 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
copybara-service bot pushed a commit to google-deepmind/clrs that referenced this pull request Oct 2, 2023
… 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
copybara-service bot pushed a commit to google-deepmind/clrs that referenced this pull request Oct 2, 2023
… 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
copybara-service bot pushed a commit to google-research/scenic that referenced this pull request Oct 2, 2023
… 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
copybara-service bot pushed a commit to google-research/scenic that referenced this pull request Oct 2, 2023
… 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
copybara-service bot pushed a commit to google-research/scenic that referenced this pull request Oct 2, 2023
… 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
copybara-service bot pushed a commit to google-research/scenic that referenced this pull request Oct 2, 2023
… 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
@copybara-service copybara-service bot merged commit c3e73c6 into jax-ml:main Oct 3, 2023
@superbobry superbobry deleted the array-any branch October 4, 2023 11:09
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
pull ready Ready for copybara import and testing
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants