Implement `StructuredDotGradCSR` and `StructuredDotGradCSC` in numba backend by tomicapretto · Pull Request #1860 · pymc-devs/pytensor

tomicapretto · 2026-01-29T03:03:37Z

Description

The main contribution of this PR is the implementation of StructuredDotGradCSR and StructuredDotGradCSC in the numba backend.

While I was working on it, I noticed Ops SpSum and SparseFromDense were running in object mode, so I also implemented them.

Checklist

Checked that the pre-commit linting/style checks pass
Included tests that prove the fix is effective or that the new feature works
Added necessary documentation (docstrings and/or example notebooks)
If you are a pro: each commit corresponds to a relevant logical change

Type of change

pytensor/link/numba/dispatch/sparse/math.py

pytensor/link/numba/dispatch/sparse/variable.py

pytensor/link/numba/dispatch/sparse/math.py

tomicapretto · 2026-01-31T18:26:32Z

The test that fails is:

FAILED tests/tensor/test_slinalg.py::TestSchur::test_schur_empty - ValueError: negative dimensions not allowed

which is unrelated to this PR.

ricardoV94 · 2026-01-31T21:23:32Z

The test that fails is:
FAILED tests/tensor/test_slinalg.py::TestSchur::test_schur_empty - ValueError: negative dimensions not allowed
which is unrelated to this PR.

@jessegrabowski

pytensor/link/numba/dispatch/sparse/basic.py

ricardoV94 · 2026-01-31T21:29:49Z

pytensor/link/numba/dispatch/sparse/math.py

+
+                for col_idx in range(size):
+                    for value_idx in range(x_ptr[col_idx], x_ptr[col_idx + 1]):
+                        output[value_idx] = np.dot(


Have to be careful with np.dot. IIRC numba overload doesn't support integer / mixed dtypes well

Argh, I'm using it since np.sum(x * y) was slower. There are a bunch of test that pass different data types, and they have all passed. Probably that's ok?

Its probably fine as long as we're upcasting the inputs to a common dtype in the make_node of Dot?

In the medium term we should consider re-implementing the BLAS calls ourselves

Do we have mixed integer / float types in the test. Or just discrete. I have >20% belief numba np.dot overload refuses those

You're right. The issue was not covered by tests (I mixed something I saw in tests for Dot with what we're doing here with StructuredDot) plus Numba does not accept mixed types in np.dot.

Rejected as the implementation raised a specific error: TypingError: np.dot() arguments must all have the same dtype

I'll double check the upcasting

This doesn't solve the numba side? In numba you'll have to explicitly cast one or both of the inputs to dot them?

You're so right!

What do you think about this?

https://github.com/tomicapretto/pytensor/blob/575ecebc6c093ef23d7d6b966140ecdc15803306/pytensor/link/numba/dispatch/sparse/math.py#L509-L522

Probably fine, but perhaps add a print statement under a if pytensor.config.compiler_verbose so users can track down a potential source of slow down. We do this for a couple of linalg numba dispatch.

You should be able to know at dispatch time whether a conversion will be needed (saying this because the warning shouldn't be inside the numba function)

Just added it, thanks for the hint @ricardoV94

ricardoV94

This looks great, I just left some minor comments

jessegrabowski · 2026-02-01T15:48:25Z

pytensor/link/numba/dispatch/sparse/math.py

+    axis = op.axis
+
+    @numba_basic.numba_njit
+    def perform(x):


does mypy freak out if you typehint this as SparseArray -> SparseArray? It would make the function more clear. Not required if it causes a headache (typehinting these overloads often does)

I have not tried it, but the SpSum op returns a dense array (see this).

What happens here is that this calls the function I implemented in overload_sum in variable.py.
Maybe a global somewhere (per op or at the top of the file) saying that many (if not all) Ops are using overloads written in a separate python file?

jessegrabowski · 2026-02-01T15:57:24Z

pytensor/link/numba/dispatch/sparse/math.py

        # General spmspm algorithm in CSR format
        @numba_basic.numba_njit
-        def _spmspm(n_row, n_col, x_ptr, x_ind, x_data, y_ptr, y_ind, y_data):
+        def _spmspm_csr(x, y, n_row, n_col):


I think it's worth considering a bit of reorganization here for future extensibility. We can make a new sparse/math sub-module and have a sum.py file with each of these inner njit functions defined independently. numba_funcify_SparseDenseMultiply can still live here, but it would be just an input checker and routing to the correct function. I'm thinking about what it will look like in the future to add support for a new sparse type.

The pattern I'm thinking about is what we are doing with linalg, for example QZ: each case is defined separately here, then the actual dispatch is defined here.

It sounds good to me. I thought a bit about it prior starting to work on this, but I saw the other ops in this module were implemented this way, so I thought it was for a reason. Maybe I just overthought about it and it was simple convenience.

pytensor/link/numba/dispatch/sparse/math.py

pytensor/link/numba/dispatch/sparse/variable.py

pytensor/sparse/math.py

tests/link/numba/sparse/test_math.py

pytensor/link/numba/dispatch/sparse/variable.py

… backend

…those implementations in SparseFromDense

ricardoV94 · 2026-02-05T12:19:52Z

Is this ready for final review?

tomicapretto · 2026-02-05T13:29:31Z

Is this ready for final review?

Yes

tomicapretto force-pushed the sparse_gradients_numba branch 2 times, most recently from 190c587 to 4690cde Compare January 29, 2026 13:34

ricardoV94 reviewed Jan 29, 2026

View reviewed changes

pytensor/link/numba/dispatch/sparse/math.py Outdated Show resolved Hide resolved

tomicapretto force-pushed the sparse_gradients_numba branch 2 times, most recently from c96ae8c to 2025883 Compare January 29, 2026 14:56

tomicapretto commented Jan 29, 2026

View reviewed changes

pytensor/link/numba/dispatch/sparse/variable.py Show resolved Hide resolved

tomicapretto force-pushed the sparse_gradients_numba branch 2 times, most recently from 512fb59 to 32099f1 Compare January 30, 2026 03:42

tomicapretto commented Jan 30, 2026

View reviewed changes

pytensor/link/numba/dispatch/sparse/math.py Show resolved Hide resolved

tomicapretto commented Jan 30, 2026

View reviewed changes

pytensor/link/numba/dispatch/sparse/math.py Outdated Show resolved Hide resolved

tomicapretto commented Jan 30, 2026

View reviewed changes

pytensor/link/numba/dispatch/sparse/math.py Show resolved Hide resolved

tomicapretto commented Jan 30, 2026

View reviewed changes

pytensor/link/numba/dispatch/sparse/math.py Show resolved Hide resolved

tomicapretto force-pushed the sparse_gradients_numba branch 3 times, most recently from a054b5d to 6af5a1a Compare January 31, 2026 17:54

tomicapretto marked this pull request as ready for review January 31, 2026 17:56

tomicapretto added enhancement New feature or request numba sparse variables labels Jan 31, 2026

ricardoV94 reviewed Jan 31, 2026

View reviewed changes

pytensor/link/numba/dispatch/sparse/basic.py Outdated Show resolved Hide resolved

ricardoV94 reviewed Jan 31, 2026

View reviewed changes

pytensor/link/numba/dispatch/sparse/basic.py Show resolved Hide resolved

ricardoV94 reviewed Jan 31, 2026

View reviewed changes

jessegrabowski reviewed Feb 1, 2026

View reviewed changes

tomicapretto marked this pull request as draft February 2, 2026 13:24

tomicapretto force-pushed the sparse_gradients_numba branch 2 times, most recently from e19f795 to 38e92e1 Compare February 2, 2026 13:40

ricardoV94 reviewed Feb 2, 2026

View reviewed changes

pytensor/link/numba/dispatch/sparse/variable.py Outdated Show resolved Hide resolved

tomicapretto force-pushed the sparse_gradients_numba branch 3 times, most recently from b886acd to 63775d6 Compare February 2, 2026 14:18

tomicapretto mentioned this pull request Feb 2, 2026

Reorganize pytensor/link/numba/dispatch/sparse/math.py #1874

Open

tomicapretto marked this pull request as ready for review February 2, 2026 14:29

tomicapretto force-pushed the sparse_gradients_numba branch 3 times, most recently from a2e9294 to 575eceb Compare February 3, 2026 00:59

tomicapretto added 5 commits February 5, 2026 09:04

Implement StructuredDotGradCSR and StructuredDotGradCSC in numba backend

026202a

Implement SparseFromDense in numba backend

59f0217

Implement SpSum in numba backend

de647d0

Add tests for: SparseFromDense, StructuredDotGrad, and SpSum in numba…

33b10f9

… backend

Extend the implementation of csr_matrix and csc_matrix in numba. Use …

2873849

…those implementations in SparseFromDense

tomicapretto force-pushed the sparse_gradients_numba branch from 575eceb to 2873849 Compare February 5, 2026 12:04

Conversation

tomicapretto commented Jan 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Checklist

Type of change

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tomicapretto commented Jan 31, 2026

Uh oh!

ricardoV94 commented Jan 31, 2026

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ricardoV94 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ricardoV94 commented Feb 5, 2026

Uh oh!

tomicapretto commented Feb 5, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

tomicapretto commented Jan 29, 2026 •

edited

Loading