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

Add underlying method argument to jax.numpy.digitize #23669

Merged
merged 1 commit into from
Sep 17, 2024

Conversation

tchatow
Copy link
Contributor

@tchatow tchatow commented Sep 16, 2024

Fixes #22489 by providing a means to choose an underlying searchsorted implementation compatible with shape poly.

jax/_src/numpy/lax_numpy.py Outdated Show resolved Hide resolved
tests/lax_numpy_test.py Show resolved Hide resolved
tests/lax_metal_test.py Outdated Show resolved Hide resolved
@jakevdp jakevdp self-assigned this Sep 16, 2024
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, thanks!

@google-ml-butler google-ml-butler bot added kokoro:force-run pull ready Ready for copybara import and testing labels Sep 17, 2024
@copybara-service copybara-service bot merged commit efa270a into jax-ml:main Sep 17, 2024
12 of 14 checks passed
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.

jax.numpy.digitize doesn't work with shape polymorphism
3 participants