Numba CAReduce: respect acc_dtype #1773
Merged
+190
−83
Conversation
This file contains hidden or 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
force-pushed
the
numba_careduce
branch
2 times, most recently
from
3c61b88 to
ad960d6
Compare
There was a problem hiding this comment.
Pull request overview
This PR fixes the Numba backend implementation of CAReduce operations to properly respect the acc_dtype (accumulation dtype) parameter and corrects the handling of infinity identities for unsigned integer types.
Key Changes:
- Modified
create_multiaxis_reducerto properly handleacc_dtypeparameter separate fromout_dtype, including support for complex-to-real conversions - Fixed infinity identity values for discrete (unsigned/signed integer) types by replacing infinite values with appropriate min/max values
- Updated
__str__methods inCAReduceandFixedOpCAReduceto display accumulation dtype when it differs from output dtype
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
| tests/link/numba/test_elemwise.py | Adds three new tests: verifying acc_dtype is respected during accumulation, testing complex-to-float conversions, and testing discrete infinity identities |
| pytensor/tensor/math.py | Updates FixedOpCAReduce.__str__() to display acc_dtype in string representation when different from output dtype |
| pytensor/tensor/elemwise.py | Updates CAReduce.__str__() to display acc_dtype in string representation when different from output dtype |
| pytensor/link/numba/dispatch/elemwise.py | Core implementation changes: refactors create_multiaxis_reducer to handle acc_dtype parameter, fixes infinity identities for discrete types, updates Softmax/SoftmaxGrad/LogSoftmax to use keyword arguments, adds cache versioning |
Also fix infinity identities for unsigned integers
ricardoV94
force-pushed
the
numba_careduce
branch
from
ad960d6 to
690015f
Compare
jessegrabowski
approved these changes
Dec 13, 2025
| res_shape = x_shape[2] | ||
| res = np.full(res_shape, numba_basic.to_scalar(0.0), dtype=out_dtype) | ||
| res_shape = (x_shape[0], x_shape[1]) | ||
| # identity = 0.0 |
Member
Author
There was a problem hiding this comment.
It's to make codegen readable when debugging, otherwise you just see a global identity being used but won't know its value
Comment on lines
205
to
206
| # if len(axes) == 1: | ||
| # return create_axis_reducer(scalar_op, identity, axes[0], ndim, dtype) |
Member
There was a problem hiding this comment.
Might as well remove this while we're here?
| def {careduce_fn_name}(x): | ||
| x_shape = x.shape | ||
| res_shape = {res_shape} | ||
| # identity = {identity} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
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.
Also fix infinity identities for unsigned integers
Cherry pick from #811