Add some sparse ops and rename others to have more explicit names by tomicapretto · Pull Request #2469 · pymc-devs/pytensor

tomicapretto · 2026-10-09T13:12:02Z

Description

This PR adds two Ops for sparse matrices, implements those and a few others in Numba, and gives more explicit names to several Ops.

Renames

Renamed all Ops that contained names such as AddSS or AddSD to be more explicit, now they are AddSparseSparse and AddSparseDense, respectively.

Implementation

Add HasCanonicalFormat, a new Op for sparse matrices that determines whether indices are unique and sorted. It is used to then implement the has_canonical_format attribute in _sparse_py_operators which is inherited by sparse matrices in PyTensor. This attribute is used in the Numba backend to compute the sum between two sparse mamtrices. However, I am not 100% sure a new Op should be added for this need in one of the backends. Thus I'm raising the question here.
Add SumDuplicates, a new Op for sparse matrices that eliminates entries by adding them together. It is not used when the matrix is already in the canonical format. In SciPy this is an in-place operation, but we return a new one. This Op is used in a few places in the Numba backend to support the implementations of the Ops mentioned below.
Numba implementations: AddSparseSparse, MultiplySparseSparse, AddSparseDense, AddSparseSparseData, StructuredAddSparseDenseVector.

Comment about numba implementations: I've discovered that one can get more performance if:

Variables used to index arrays are marked as unsigned.
Casts are performed at the scalar level, not at the array level (which requires a larger allocation).

In a following PR (or this one?) I can incorporate these ideas in the previously implemented ops.

Other

Updated test_sparse_creation_refcount and test_sparse_passthrough_refcount in tests/link/numba/sparse/test_basic.py to use the effective initial reference count, which can differ between Python versions. Right now it's hardcoded to 3, but it is 2 in Python 3.14+.

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

ricardoV94 · 2026-10-10T15:48:28Z


+    @pytest.mark.parametrize("format", ["csr", "csc"])
+    def test_inplace(self, format):
+        x = getattr(sparse, format + "_matrix")("x", dtype="float64")


why don't you pass it directly in the parametrize, instead of retrieving by getattr?

It's updated now

ricardoV94 · 2026-10-10T15:48:32Z

        self.inplace = inplace
        if self.inplace:
-            self.view_map = {0: [0]}
+            self.destroy_map = {0: [0]}


But is it a view when not inplace?

Nope, .sorted_indices() returns a copy https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.csr_matrix.sorted_indices.html.

then it's never inplace / aliased either?

When inplace, it uses .sort_indices(), which mutates x inplace, and it then returns x. When not inplace it used sorted_indices(), which returns a copy. It's like list.sort and sorted(list), right?

omg sort vs sorted is the difference. ok

ricardoV94 · 2026-10-10T15:49:03Z

        self.op = ensure_sorted_indices

+    @pytest.mark.parametrize("format", ["csr", "csc"])
+    def test_inplace(self, format):


how is this testing inplace (mutation of the original input)?

I've updated it now

ricardoV94 · 2026-10-10T15:50:16Z

                verify_grad_sparse(self.op, data, structured=False)


+class TestSumDuplicates(utt.InferShapeTester):


missing a perform test (gradient tests perform, but can't validate it by itself?)

ricardoV94 · 2026-10-14T13:09:09Z

can you add a summary of this PR in the top message. It seems it's not just renaming but creating new Ops as well?

tomicapretto · 2026-10-14T13:11:54Z

can you add a summary of this PR in the top message. It seems it's not just renaming but creating new Ops as well?

Yes, but it's not done yet (that's why I am keeping it as a draft). I'll add a summary and mark as ready once it's good for review. But of course thanks for the feedback already given!

…ving explicit zeros.

…parseSparseData, and AddStructuredSparseDenseVector in Numba

tomicapretto force-pushed the sparse-names branch from 3fe0ca2 to b7c4de9 Compare October 9, 2026 14:12

ricardoV94 reviewed Oct 10, 2026

View reviewed changes

tomicapretto force-pushed the sparse-names branch 2 times, most recently from 8de2d0d to 9c63961 Compare October 14, 2026 13:02

tomicapretto changed the title ~~Rename sparse ops to have more explicit names~~ Implement some sparse ops and others to have more explicit names Oct 14, 2026

tomicapretto changed the title ~~Implement some sparse ops and others to have more explicit names~~ Add some sparse ops and rename others to have more explicit names Oct 14, 2026

tomicapretto force-pushed the sparse-names branch 4 times, most recently from f3d3bd5 to 9552550 Compare October 16, 2026 17:35

tomicapretto added 9 commits October 16, 2026 17:33

Rename sparse ops to have more explicit names

bec5222

Fix mutation tracking for inplace EnsureSortedIndices

59fff28

Add SumDuplicates Op to combine duplicate sparse entries while preser…

d0e6a04

…ving explicit zeros.

Implement EnsureSortedIndices in numba

cec8049

Add HasCanonicalFormat Op for sparse matrices

a4dd2ad

Implement HasCanonicalFormat in numba

b5fe501

Implement SumDuplicates in numba

f61ad6c

Implement AddSparseSparse, MultiplySparseSparse, AddSparseDense, AddS…

8782c8b

…parseSparseData, and AddStructuredSparseDenseVector in Numba

Adapt sparse data test to use original refcount

c79b358

tomicapretto force-pushed the sparse-names branch from 9552550 to c79b358 Compare October 16, 2026 21:05

tomicapretto marked this pull request as ready for review October 16, 2026 21:08

tomicapretto requested a review from ricardoV94 October 16, 2026 21:08

tomicapretto mentioned this pull request Oct 17, 2026

Improve performance of several Ops in sparse numba #2505

Open

6 tasks done

		verify_grad_sparse(self.op, data, structured=False)


		class TestSumDuplicates(utt.InferShapeTester):

Conversation

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

Uh oh!

Description

Renames

Implementation

Other

Checklist

Type of change

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 Oct 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

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

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 commented Oct 14, 2026

Uh oh!

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

Uh oh!

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

tomicapretto commented Oct 9, 2026 •

edited

Loading

ricardoV94 Oct 14, 2026 •

edited

Loading

tomicapretto Oct 14, 2026 •

edited

Loading

tomicapretto commented Oct 14, 2026 •

edited

Loading