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

jax.numpy: implement scalar boolean indexing #19722

Merged
merged 1 commit into from
Feb 9, 2024

Conversation

jakevdp
Copy link
Collaborator

@jakevdp jakevdp commented Feb 8, 2024

Part of #18353

@jakevdp jakevdp added the enhancement New feature or request label Feb 8, 2024
@jakevdp jakevdp self-assigned this Feb 8, 2024
@jakevdp jakevdp mentioned this pull request Feb 8, 2024
38 tasks
@jakevdp jakevdp force-pushed the scalar-bool-indexing branch 3 times, most recently from 41bc918 to 0f68834 Compare February 8, 2024 22:12
@jakevdp jakevdp requested a review from froystig February 8, 2024 22:21
raise IndexError(
f"Too many indices for array: array has ndim of {self.ndim}, but "
f"was indexed with {num_idx} non-None/Ellipsis indices.")

Copy link
Collaborator Author

@jakevdp jakevdp Feb 8, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Removed this because the logic is wrong (boolean indices idx consume a number of dimensions equal to idx.ndim) and this condition is correctly checked in _index_to_gather.

Copy link
Member

@froystig froystig left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nicely done!

@google-ml-butler google-ml-butler bot added kokoro:force-run pull ready Ready for copybara import and testing labels Feb 9, 2024
@copybara-service copybara-service bot merged commit 3d7ed94 into jax-ml:main Feb 9, 2024
14 checks passed
@jakevdp jakevdp deleted the scalar-bool-indexing branch February 9, 2024 19:43
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request pull ready Ready for copybara import and testing
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants