Skip to content

ENH: at support for bool mask in Dask and JAX #121

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

Merged
merged 4 commits into from
Jan 25, 2025

Conversation

crusaderky
Copy link
Contributor

@crusaderky crusaderky commented Jan 22, 2025

This special case treatment should be incorporated inside jax itself.
In the meantime, for the sake of expedience, I'm implementing it here.

@crusaderky crusaderky changed the title WIP ENH: at support for bool mask in Dask and JAX ENH: at support for bool mask in Dask and JAX Jan 22, 2025
@lucascolley lucascolley added lazy arrays enhancement New feature or request labels Jan 22, 2025
@crusaderky crusaderky force-pushed the at_bool_mask branch 2 times, most recently from 4863400 to 6bb7d96 Compare January 24, 2025 16:29

You should instead use::

>>> x = xp.where(m(x), f(x), x)
Copy link
Contributor Author

Choose a reason for hiding this comment

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

TODO suggest to use lazywhere

@crusaderky crusaderky marked this pull request as ready for review January 24, 2025 19:37
@crusaderky
Copy link
Contributor Author

@lucascolley ready for review!

@lucascolley lucascolley self-requested a review January 24, 2025 19:43
crusaderky and others added 2 commits January 25, 2025 13:40
Co-authored-by: Lucas Colley <lucas.colley8@gmail.com>
Copy link
Member

@lucascolley lucascolley left a comment

Choose a reason for hiding this comment

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

thanks, LGTM!

@lucascolley lucascolley added this to the 0.6.1 milestone Jan 25, 2025
@lucascolley lucascolley merged commit a5cb116 into data-apis:main Jan 25, 2025
10 checks passed
@crusaderky crusaderky deleted the at_bool_mask branch January 25, 2025 23:42
NeilGirdhar pushed a commit to NeilGirdhar/array-api-extra that referenced this pull request Feb 20, 2025
Co-authored-by: Lucas Colley <lucas.colley8@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request lazy arrays
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants