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

WIP jax.numpy array API compliance finalization #21323

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from

Conversation

Micky774
Copy link
Collaborator

@Micky774 Micky774 commented May 21, 2024

Towards #21088

Changes

  • Adds the following attributes to ArrayImpl:
    • __array_namespace__ property
    • to_device method
    • device property
  • Adds the following to jax.numpy:
    • __array_namespace_info__
    • __array_api_version__

Notes

This PR is a draft right now since we should include these changes last so as to publicly support jax.numpy as an array API compliant namespace. @vfdev-5 can take over this PR later once the remainder of the work is completed.

This does not need to wait on all the ongoing array API related deprecations to be completed, since some of them are only required for the 2023 standard, hence we can likely adopt the 2022 standard first.

It may make sense break off the to_device and device changes for ArrayImpl into a small separate PR, since they don't imply explicit compliance by themselves, but I wanted to keep them together in this PR in case there were any caveats wrt to ArrayImpl vs Tracer behaviors that we should discuss first (based on old TODO note).

@jakevdp
Copy link
Collaborator

jakevdp commented May 21, 2024

We only finally removed the arr.device() method in JAX v0.4.27 – to avoid confusion for users I think we should wait for one more release (0.4.29) before we add the arr.device property, so that it will be part of 0.4.30. What do you think?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

2 participants