Speed up MatrixExpr.add.reduce via quicksum
#1157
Open
+128
−76
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
Implements the __array_ufunc__ method in MatrixExpr to handle numpy ufuncs, specifically enabling correct behavior for reductions like np.add.reduce by delegating to the sum method. This improves compatibility with numpy operations.
Moved the sum computation logic from MatrixExpr.sum and __array_ufunc__ into a new _core_sum function for better code reuse and maintainability.
Introduces tests to compare the performance of the mean operation on matrix variables and checks the return types of mean with and without axis argument.
Documented the speed improvement for MatrixExpr.add.reduce using quicksum in the changelog.
The sum method was removed from the MatrixExpr class, consolidating summation logic in the _core_sum function. The docstring for _core_sum was expanded to include detailed parameter and return value descriptions, improving code clarity and maintainability.
Zeroto521
commented
Jan 16, 2026
| def __array_ufunc__(self, ufunc, method, *args, **kwargs): | ||
| if method == "reduce": | ||
| if ufunc is np.add and isinstance(args[0], MatrixExpr): | ||
| return _core_sum(args[0], **kwargs) |
Contributor
Author
There was a problem hiding this comment.
It can't use MatrixExpr.sum directly.
Or it will crush the Python kernel.
Use Matrix.sum the calling chain like:
x.mean(1)x.sum(1)np.add.reduce__array_ufunc__x.sum(1)__array_ufunc__np.add.reduce- ...
Deleted the type stub for the sum method in MatrixExpr, likely to reflect changes in the underlying implementation or to correct type information.
Enhanced the __array_ufunc__ method in MatrixExpr to ensure proper array conversion and consistent return types. Added the _ensure_array helper for argument handling. Also removed an unused include of matrix.pxi from expr.pxi.
Zeroto521
commented
Jan 16, 2026
| # which should, in princple, modify the expr. However, since we do not implement __isub__, __sub__ | ||
| # gets called (I guess) and so a copy is returned. | ||
| # Modifying the expression directly would be a bug, given that the expression might be re-used by the user. </pre> | ||
| include "matrix.pxi" |
Contributor
Author
There was a problem hiding this comment.
scip.pyi also has this line.
cdef function (_ensure_array) can't be added twice.
Updated MatrixExpr.__matmul__ to return the correct type when the result is not an ndarray. Adjusted tests to reflect the expected return type for 1D matrix multiplication and improved performance test timing logic.
Simplifies timing measurement in matrix sum performance tests by directly calculating elapsed time instead of storing start and end times separately. This improves code readability and reduces variable usage.
Changed the expected type of 1D @ 1D matrix multiplication from MatrixExpr to Expr in test_matrix_matmul_return_type to reflect updated behavior.
Updated tests to use x.view(MatrixExpr) and x.view(np.ndarray) instead of direct subclass method calls. This clarifies the intent and ensures the correct method resolution for sum and mean operations in performance and result comparison tests.
Zeroto521
commented
Jan 16, 2026
| end_orig = time() | ||
| # Original sum via `np.ndarray.sum` | ||
| start = time() | ||
| x.view(np.ndarray).sum() |
Contributor
Author
There was a problem hiding this comment.
np.ndarray.sum is MatrixExpr.sum now. They will both call __array_ufunc__ protocol. So it has to convert MatrixVariable to a normal np.ndarray.
Zeroto521
commented
Jan 16, 2026
|
|
||
| np_res = a.sum(axis, keepdims=keepdims) | ||
| scip_res = MatrixExpr.sum(a, axis, keepdims=keepdims) | ||
| scip_res = a.view(MatrixExpr).sum(axis, keepdims=keepdims) |
Contributor
Author
There was a problem hiding this comment.
np.ndarray.sum is MatrixExpr.sum now.
MatrixExpr.sum(a) will return np.ndarray(..., dtype=int) not MatrixExpr
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
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.
closes to #1153
np.add.reduceis the inner function fornp.ndarray.sum,np.ndarray.mean, and more.