Reorganize the sparse module

[jessegrabowski]

Description

Reorganize the existing sparse code to be a bit more logical. In particular, the sparse module now mirrors the tensor module:

• sparse/basic.py has constructors and utilities
• sparse/math.py has math stuff
• sparse/variable.py has base classes
• sparse/linalg.py has linear algebra

Other stuff was mostly untouched. My motivation here is to lay the groundwork to get people to come look at and work on sparse. Things were pretty messy before. There was a lot of magic I didn't like. For example, the _sparse_py_operators inherited from _tensor_py_operators, and did some wrapper magic to dispatch a huge list of methods to sparse. I made this a lot more verbose, and put NotImplemented everywhere, so it's obvious to potential helpers what needs to be implemented.

Related Issue

• Closes #
• Related to #

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

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

---

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

[ricardoV94]

I think it was a mistake to inherit from TensorVariable in the first place. They are way too distinct objects

[ricardoV94]

I believe we should revert aesara-devs/aesara#142, in fact the example given of Subtensor just further proves the point. Indexing in sparsevariables has nothing to do with in indexing in tensorvariables, starting with the fact there are no sparse vectors, only matrices

[jessegrabowski]

mypy is failing because of that inheritance for what its worth

[jessegrabowski]

I believe we should revert aesara-devs/aesara#142, in fact the example given of Subtensor just further proves the point. Indexing in sparsevariables has nothing to do with in indexing in tensorvariables, starting with the fact there are no sparse vectors, only matrices

I agree. The current mixins don't make any sense at all. We have e.g. SparseShared inheriting from two different sources that have the same method twice. So now the order matters, which sucks. I opened an issue to address it. For this PR, I'm just having mypy ignore it.

[jessegrabowski]

This leads to 2,304 test cases, which seems slightly excessive.

[ricardoV94]

only?

[jessegrabowski]

It was marked as a slow test, which isn't super shocking. But it actually runs in <1s on my local machine after all your recent compile speedups. Pretty impressive!

[ricardoV94]

Could be a caching thing. Our CI doesn't have that luxury, so it may be slow in the CI still (we can check)

[ricardoV94]

it takes 60s without cache on my machine, and 30s on the second run. I don't know how you get <1s, that seems wild, perhaps you're just seeing the time of "skipping"?

[codecov]

Codecov Report

❌ Patch coverage is 81.99234% with 235 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.69%. Comparing base (1f9a67b) to head (c904d24).

Files with missing lines	Patch %	Lines
pytensor/sparse/math.py	81.79%	88 Missing and 70 partials ⚠️
pytensor/sparse/variable.py	84.47%	42 Missing and 1 partial ⚠️
pytensor/sparse/basic.py	78.30%	13 Missing and 10 partials ⚠️
pytensor/sparse/rewriting.py	65.38%	2 Missing and 7 partials ⚠️
pytensor/sparse/linalg.py	91.66%	1 Missing and 1 partial ⚠️

❌ Your patch check has failed because the patch coverage (81.99%) 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    #1674      +/-   ##
==========================================
- Coverage   81.70%   81.69%   -0.02%     
==========================================
  Files         246      249       +3     
  Lines       53632    53824     +192     
  Branches     9438     9436       -2     
==========================================
+ Hits        43820    43971     +151     
- Misses       7330     7372      +42     
+ Partials     2482     2481       -1     

Files with missing lines	Coverage Δ
pytensor/link/jax/dispatch/sparse.py	83.33% <100.00%> (ø)
pytensor/sparse/sharedvar.py	94.11% <100.00%> (ø)
pytensor/tensor/slinalg.py	91.40% <100.00%> (ø)
pytensor/sparse/linalg.py	91.66% <91.66%> (ø)
pytensor/sparse/rewriting.py	75.44% <65.38%> (+0.09%)	⬆️
pytensor/sparse/basic.py	82.86% <78.30%> (+0.17%)	⬆️
pytensor/sparse/variable.py	84.47% <84.47%> (ø)
pytensor/sparse/math.py	81.79% <81.79%> (ø)

... and 2 files with indirect coverage changes

🚀 New features to boost your workflow:

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

[jessegrabowski]

@ricardoV94 look at this one too since you're on a mergeathon

[ricardoV94]

I'm more comfortable letting @Armavica review this one. He somehow has an eye for walls of linediffs :)

[ricardoV94]

How confident are you that you only did a refactor and nothing was lost.

Also is this a major change, or nothing changes in user facing API (like import locations, Op names)? Need to know for the PR labels

[ricardoV94]

It seems like you changed somethings. .exp() now returns structured_exp(self) but that's wrong. I don't think structured_x should exist for anything that doesn't map 0 -> 0, but that's another story. It may have happened to a few more