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.

Sign up for GitHub

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

Jump to bottom

[array API] add device property & to_device method #22597

Merged
merged 1 commit into from
Jul 23, 2024
Merged

[array API] add device property & to_device method #22597

merged 1 commit into from
Jul 23, 2024

Conversation

jakevdp
Copy link
Collaborator

@jakevdp jakevdp commented Jul 23, 2024 •
edited
Loading

Part of #21088 and #18353

As discussed earlier, arr.device will return the device for single-device arrays, and the sharding for non-trivially sharded arrays. This lets arr.device be passed to device_put, as well as to other Array API arguments that accept Device | Sharding.

@jakevdp jakevdp self-assigned this Jul 23, 2024
@jakevdp jakevdp requested a review from yashk2810 July 23, 2024 15:22
yashk2810
yashk2810 reviewed Jul 23, 2024
View reviewed changes
jax/_src/array.py Show resolved Hide resolved
jax/_src/earray.py Show resolved Hide resolved
jax/_src/numpy/array_methods.py Outdated Show resolved Hide resolved
@jakevdp jakevdp requested a review from yashk2810 July 23, 2024 15:37
@google-ml-butler google-ml-butler bot added kokoro:force-run pull ready Ready for copybara import and testing labels Jul 23, 2024
copybara-service bot pushed a commit that referenced this pull request Jul 23, 2024
@ezhulenev
[xla:cpu] Support for up to 16 sorted inputs
b98c1e0 
+ enable more jax/lax tests for XLA CPU thunks

FUTURE_COPYBARA_INTEGRATE_REVIEW=#22597 from jakevdp:arr-device 613a000
PiperOrigin-RevId: 654865806
@copybara-service copybara-service bot mentioned this pull request Jul 23, 2024
Merged
@copybara-service copybara-service bot merged commit dc42ba0 into jax-ml:main Jul 23, 2024
12 of 15 checks passed
@jakevdp jakevdp deleted the arr-device branch July 23, 2024 18:28
@jakevdp jakevdp mentioned this pull request Jul 29, 2024
Closed
38 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@yashk2810 yashk2810 yashk2810 approved these changes

Assignees

@jakevdp jakevdp

Labels
pull ready Ready for copybara import and testing
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@jakevdp @yashk2810 @kokoro-team