pytensor
pytensor copied to clipboard
Optimize non-structured unary elementwise `Ops` on sparse matrices
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.