Add Numba implementation of Blockwise #1015
Open
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
ricardoV94
added
enhancement
ricardoV94
force-pushed
the
numba_blockwise
branch
from
8f514ac
to
2720d0b
Compare
ricardoV94
force-pushed
the
numba_blockwise
branch
2 times, most recently
from
2014cd9
to
c45aab2
Compare
ricardoV94
commented
Oct 6, 2024
Comment on lines
+62
to
+74
| if nout == 1: | ||
| tuple_core_shapes = (to_fixed_tuple(core_shapes[0], core_shape_0),) | ||
| elif nout == 2: | ||
| tuple_core_shapes = ( | ||
| to_fixed_tuple(core_shapes[0], core_shape_0), | ||
| to_fixed_tuple(core_shapes[1], core_shape_1), | ||
| ) | ||
| else: | ||
| tuple_core_shapes = ( | ||
| to_fixed_tuple(core_shapes[0], core_shape_0), | ||
| to_fixed_tuple(core_shapes[1], core_shape_1), | ||
| to_fixed_tuple(core_shapes[2], core_shape_2), | ||
| ) |
There was a problem hiding this comment.
If anybody has an idea on how to do this dynamically would be great. Do we have to do string generation 😭?
There was a problem hiding this comment.
Are you opposed to cheesing it?
tuple(to_fixed_tuple(core_shapes[i], core_shape_lens[i]) for i in range(nout))(I don't have full context)
There was a problem hiding this comment.
numba doesn't support that in this context
There was a problem hiding this comment.
Ewww. Maybe you could try a bunch of eval statements?
There was a problem hiding this comment.
We would need to go down the string generation as we do for some other Ops (like Scan). But I didn't want to :)
This can only be done when the output of infer_shape of the core_op depends only on the input shapes, and not their values.
Restricted to 3 outputs, due to limitations in jitting of Numba functions
ricardoV94
force-pushed
the
numba_blockwise
branch
from
c45aab2
to
31cc1e9
Compare
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1015 +/- ##
==========================================
- Coverage 81.75% 81.74% -0.01%
==========================================
Files 183 185 +2
Lines 47756 47816 +60
Branches 11620 11632 +12
==========================================
+ Hits 39044 39089 +45
- Misses 6519 6529 +10
- Partials 2193 2198 +5
|
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Description
Implement Numba blockwise for Ops with up to 3 outputs (due to numba not liking tuple generators in the inner functions...)
It uses the machinery developed for RVs and Elemwise. The hard part has to do with multiple number of inputs and numba fussiness.
It also improves Blockwise shape inference based on the infer_shape of the core ops
The small cholesky benchmark I added here test runs 10x faster after this PR on my local machine.
Related Issue