pytensor icon indicating copy to clipboard operation
pytensor copied to clipboard

Optimize non-structured unary elementwise `Ops` on sparse matrices

Open jessegrabowski opened this issue 1 month ago • 3 comments

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.

jessegrabowski avatar Oct 18 '25 18:10 jessegrabowski