pymc-devs / pymc-devs/pytensor
Optimize non-structured unary elementwise `Ops` on sparse matrices
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 644
- Forks
- 208
- Avg merge
- 2d 14h
- Merged PRs (30d)
- 16
Description
Description
The title sucks, but the case I'm interested in is a graph like this:
import pytensor.sparse as pts
import pytensor.tensor as pt
x = pts.csr_dmatrix()
out = pt.exp(x)
Currently we have a notion of structured_exp and "normal exp". The structured one ignores the zeros in the matrix, so you get back another sparse with the same sparsity pattern, but the non-zero elements are now exp(x.data). In the "normal exp", we do exp(x.todense())
There are two things here. First, I think a user should never get a "non-structured" output if he does pt.exp(x). That should always return a sparse matrix. If he wants a dense output, we should make him actually do pt.exp(x.todense()).
Second, if we do get pt.exp(x.todense()), we can rewrite this to do a lot better. We only need to compute pt.exp(0) once, then allocate a full matrix of that value, then finally set subtensor the non-zero values.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by tracing PyTensor's sparse unary Ops and the existing structured versus non-structured handling for exp on sparse matrices. Check the relevant sparse operation tests, if present, and compare both requested cases: pt.exp(x) should remain sparse, while pt.exp(x.todense()) should avoid repeated work on implicit zeros. Done means the behavior and optimization are covered by tests.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- backend, performance
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100