Basic Sparse functionality in Numba #1676

ricardoV94 · 2025-10-18T17:14:37Z

We had support for boxing/unboxing Sparse objects in numba, but we couldn't do anything with them.

This PR implements the basic functionality:

CSMProperties (Op that retrieves the attributes from a CSM Matrix)
CSM (Op that rebuilds a Sparse variable from the attributes)
astype()
csr_matrix, csc_matrix sicpy constructor overloads
MULSD (as a POC)

TODO

fix conflicts (from Reorganize the sparse module #1674)
Reintroduce dtype/type on the numba Sparse Matrix

📚 Documentation preview 📚: https://pytensor--1676.org.readthedocs.build/en/1676/

codecov · 2025-11-18T16:49:17Z

Codecov Report

❌ Patch coverage is 83.20000% with 42 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.67%. Comparing base (8617558) to head (1de9697).
⚠️ Report is 30 commits behind head on main.

Files with missing lines	Patch %	Lines
pytensor/link/numba/dispatch/sparse/basic.py	82.38%	24 Missing and 7 partials ⚠️
pytensor/link/numba/dispatch/sparse/math.py	90.76%	3 Missing and 3 partials ⚠️
pytensor/link/numba/dispatch/basic.py	0.00%	3 Missing ⚠️
pytensor/tensor/type.py	33.33%	1 Missing and 1 partial ⚠️

❌ Your patch check has failed because the patch coverage (83.20%) is below the target coverage (100.00%). You can increase the patch coverage or adjust the target coverage.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1676      +/-   ##
==========================================
- Coverage   81.70%   81.67%   -0.04%     
==========================================
  Files         246      251       +5     
  Lines       53632    52549    -1083     
  Branches     9438     9271     -167     
==========================================
- Hits        43822    42919     -903     
+ Misses       7329     7258      -71     
+ Partials     2481     2372     -109

Files with missing lines	Coverage Δ
pytensor/link/numba/dispatch/compile_ops.py	`92.64% <ø> (ø)`
pytensor/link/numba/dispatch/sparse/__init__.py	`100.00% <100.00%> (ø)`
pytensor/sparse/variable.py	`76.79% <100.00%> (ø)`
pytensor/tensor/type.py	`94.20% <33.33%> (-0.41%)`	⬇️
pytensor/link/numba/dispatch/basic.py	`81.77% <0.00%> (-3.25%)`	⬇️
pytensor/link/numba/dispatch/sparse/math.py	`90.76% <90.76%> (ø)`
pytensor/link/numba/dispatch/sparse/basic.py	`82.38% <82.38%> (ø)`

... and 4 files with indirect coverage changes

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

* Handle static shape * Rename to more readable Op classes * Simplify perform

2304 -> 288

Co-authored-by: Adrian Seyboldt <aseyboldt@users.noreply.github.com> Co-authored-by: Jesse Grabowski <48652735+jessegrabowski@users.noreply.github.com>

Co-authored-by: Jesse Grabowski <48652735+jessegrabowski@users.noreply.github.com>

jessegrabowski

approved with some questions

jessegrabowski · 2026-01-17T02:16:24Z

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

+        @numba_basic.numba_njit
+        def sparse_multiply_scalar(x, y):
+            if same_dtype:
+                z = x.copy()


This can't ever be inplace?

The base Op probably doesn't have inplace optimization as I basically copied the perform method. Will double check

jessegrabowski · 2026-01-17T18:45:06Z

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

+
+
+@overload(numba_deepcopy)
+def numba_deepcopy_sparse(x):


What's deep about this?

sparse_matrix.copy() does a deepcopy just like array.copy(). But for other types like list or rng there's a difference between copy and deepcopy hence the more explicit name

jessegrabowski · 2026-01-17T18:45:42Z

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

+def numba_funcify_CSMProperties(op, node, **kwargs):
+    @numba_basic.numba_njit
+    def csm_properties(x):
+        # Reconsider this int32/int64. Scipy/base PyTensor use int32 for indices/indptr.


Are we able to just go to int64 ourselves, or do we need to wait for upstream to change?

Would need to change stuff in the pre-existing Ops so that fallback to obj mode is compatible. Would leave that for a later PR if we decide

jessegrabowski · 2026-01-18T03:20:47Z

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

+    shape_obj = c.box(typ.shape, struct_ptr.shape)
+
+    # Call scipy.sparse.cs[c|r]_matrix
+    cls_obj = c.pyapi.unserialize(c.pyapi.serialize_object(typ.instance_class))


does this line mean that we always have to come back to python during construction of a numba sparse array?

No, just at the end of the outer jitted function if there's a sparse variable in the outputs. You need this for every numba type. It's where the conversion from internal numba representation to python objects happens.

If a function only uses sparse arrays internally this isn't called.

jessegrabowski · 2026-01-18T03:38:00Z

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

+@overload(sp.sparse.csr_matrix)
+def overload_csr_matrix(arg1, shape, dtype=None):
+    if not isinstance(arg1, types.BaseAnonymousTuple) or len(arg1) != 3:
+        return None


What does it mean to return None from an overload? It fails?

Overloads work by trying all registered methods until one works. So numba will keep trying until one matches

jessegrabowski · 2026-01-18T03:39:00Z

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

+    return impl
+
+
+@overload(np.shape)


how does this interact with other overloads of np.shape? e.g. what if I import this code then call np.shape on an array in numba mode, does it still work as expected?

Yes, like your question above. When this overload returns None, numba will keep trying other overloads of np.shape until one returns something other than None

jessegrabowski · 2026-01-18T03:48:47Z

pytensor/sparse/math.py

-        out[0] = self.comparison(x, y).astype("uint8")
+        # FIXME: Scipy csc > csc outputs csr format, but make_node assumes it will be the same as inputs
+        # Casting to respect make_node, but this is very inefficient
+        # TODO: Why not go with default bool?


why not indeed

I suspect some archaic C bug. We should try to remove in a separate PR

jessegrabowski · 2026-01-18T03:49:21Z

pytensor/sparse/math.py


        inputs = [x, y]  # Need to convert? e.g. assparse
-        outputs = [psb.SparseTensorType(dtype=x.type.dtype, format=myformat)()]
+        outputs = [SparseTensorType(dtype=x.type.dtype, format=myformat)()]


not your code but i hate the name myformat

jessegrabowski · 2026-01-18T03:53:29Z

tests/link/numba/sparse/test_basic.py

+    x = sp.sparse.csr_matrix(np.eye(100))
+
+    y = test_fn(x)
+    assert y is not x and np.all(x.data == y.data) and np.all(x.indices == y.indices)


do you also need to test x.data is not y.data or is that guaranteed by the first check

It would be even better to check not np.shares_memory. I'll do that

ricardoV94 added enhancement New feature or request numba sparse variables labels Oct 18, 2025

ricardoV94 force-pushed the numba_sparse_ops branch from 913e012 to 71cc54a Compare October 18, 2025 17:21

ricardoV94 force-pushed the numba_sparse_ops branch from 71cc54a to 1de9697 Compare November 18, 2025 16:25

ricardoV94 force-pushed the numba_sparse_ops branch 2 times, most recently from 14a11e2 to 285d7ca Compare November 19, 2025 12:12

This was referenced Nov 25, 2025

Numba full backend support and required dependency #811

Merged

Reconsider handlig of constants in numba backend #1759

Open

ricardoV94 force-pushed the numba_sparse_ops branch 3 times, most recently from b0057c4 to 7c2edbc Compare January 15, 2026 17:40

ricardoV94 marked this pull request as ready for review January 15, 2026 17:40

ricardoV94 changed the title ~~Implement Sparse Ops in Numba~~ Basic Sparse functionality in Numba Jan 15, 2026

ricardoV94 force-pushed the numba_sparse_ops branch 4 times, most recently from e9c1320 to dc3e431 Compare January 16, 2026 16:18

ricardoV94 added 11 commits January 16, 2026 17:53

Fix exclude tag cxx -> cxx_only

5075dfa

Seed flaky test

90c7195

Cleanup ruff per file ignores

db84708

Allow single integer as TensorType shape

fe8b80f

Handle static shape in core sparse methods

447af64

Sparse methods: Do not use deprecated names

63c8171

SparseMultiply: Cleanup Ops

0037f84

* Handle static shape * Rename to more readable Op classes * Simplify perform

CSMProperties Op: simplify logic

b6a01cc

Sparse comparison: Fix type violation in perform method

feecbfd

Test: Reduce number of parametrizations

38e27e4

2304 -> 288

Test: Reduce number of function compilation in sparse cast type

89b4f7a

ricardoV94 and others added 4 commits January 16, 2026 17:53

Numba sparse: Implement basic functionality

5333b3e

Co-authored-by: Adrian Seyboldt <aseyboldt@users.noreply.github.com> Co-authored-by: Jesse Grabowski <48652735+jessegrabowski@users.noreply.github.com>

Numba sparse: Implement basic Ops

0e078a7

Co-authored-by: Jesse Grabowski <48652735+jessegrabowski@users.noreply.github.com>

Numba sparse: Implement SparseDenseMultiply

df200e0

Co-authored-by: Jesse Grabowski <48652735+jessegrabowski@users.noreply.github.com>

Numba sparse: Remove codebase xfails

b4cccc7

ricardoV94 force-pushed the numba_sparse_ops branch from dc3e431 to b4cccc7 Compare January 16, 2026 16:54

ricardoV94 requested review from aseyboldt and jessegrabowski January 16, 2026 16:55

jessegrabowski approved these changes Jan 18, 2026

View reviewed changes

Basic Sparse functionality in Numba #1676

Are you sure you want to change the base?

Basic Sparse functionality in Numba #1676

Conversation

ricardoV94 commented Oct 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

TODO

Uh oh!

codecov bot commented Nov 18, 2025

Codecov Report

Uh oh!

jessegrabowski 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!

ricardoV94 Jan 18, 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!

ricardoV94 Jan 18, 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!

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!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

ricardoV94 commented Oct 18, 2025 •

edited

Loading

ricardoV94 Jan 18, 2026 •

edited

Loading

ricardoV94 Jan 18, 2026 •

edited

Loading