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

Support multidimensional boolean set/inc_subtensor in Numba via rewrite #1108

Merged
merged 2 commits into from
Nov 28, 2024

Conversation

ricardoV94
Copy link
Member

@ricardoV94 ricardoV94 commented Nov 28, 2024

Follow up to #818

It's simple to support multidimensional boolean set_subtensor. The corresponding integer indexing would be slightly trickier because it requires reshaping y as well.

Although there are still cases that are not covered outside of object mode (np.newaxis, non-contiguous adv index, broadcasted integer indexing, mixed boolean/integer indexing) I think it's fine to close #772. Efforts can be directed on a per-need basis. Numpy indexing is just too wild :)


📚 Documentation preview 📚: https://pytensor--1108.org.readthedocs.build/en/1108/

Setting to false can lead to slower code on C/Numba backend which don't support np.add.at natively.
@ricardoV94 ricardoV94 force-pushed the numba_bool_adv_set_subtensor branch from f94a44c to a982af9 Compare November 28, 2024 13:49
Copy link

codecov bot commented Nov 28, 2024

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 82.12%. Comparing base (9dad122) to head (8dc6727).
Report is 10 commits behind head on main.

Additional details and impacted files

Impacted file tree graph

@@            Coverage Diff             @@
##             main    #1108      +/-   ##
==========================================
- Coverage   82.14%   82.12%   -0.03%     
==========================================
  Files         183      183              
  Lines       48104    48111       +7     
  Branches     8665     8667       +2     
==========================================
- Hits        39517    39510       -7     
- Misses       6419     6435      +16     
+ Partials     2168     2166       -2     
Files with missing lines Coverage Δ
pytensor/tensor/rewriting/subtensor.py 90.47% <100.00%> (+0.07%) ⬆️
pytensor/tensor/subtensor.py 89.54% <ø> (ø)

... and 1 file with indirect coverage changes

…'t support np.add.at natively.

Support multidimensional boolean set/inc_subtensor in Numba via rewrite
@ricardoV94 ricardoV94 force-pushed the numba_bool_adv_set_subtensor branch from a982af9 to 8dc6727 Compare November 28, 2024 14:27
Copy link

@AlexAndorra AlexAndorra 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 @ricardoV94

@AlexAndorra AlexAndorra merged commit 0824dba into pymc-devs:main Nov 28, 2024
62 checks passed
@ricardoV94 ricardoV94 added the enhancement New feature or request label Dec 3, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Handle more AdvancedSetSubtensor ops in numba
2 participants