`pytensor.sparse.dot` returns a dense matrix even when both inputs are sparse

[tomicapretto]

Description

As mentioned in #1854, PyTensor's dot product for sparse matrices always returns a dense array, even when both inputs are sparse. This behavior differs from SciPy's, which returns a sparse matrix when both inputs are sparse (using the format of the first input IIRC).

Note this only happens with pytensor.sparse.dot. The other operator, pytensor.sparse.structured_dot, does return a sparse matrix when both inputs are sparse.

I have not looked into pytensor.sparse.true_dot yet, but a skim suggests this one always return a sparse matrix.

[ricardoV94]

The whole point AFAICT is that it allows the gradient with respect to the dense input (if there's one) to be a dense dot?

I'm not sure we need this extra complexity? Shall we remove?

[ricardoV94]

Ah the notes explain what's the motive:

pytensor/pytensor/sparse/math.py

Lines 1958 to 1960 in 207b0c6

	When the operation has the form ``dot(csr_matrix, dense)``
	the gradient of this operation can be performed inplace
	by `UsmmCscDense`. This leads to significant speed-ups.

[ricardoV94]

I suspect we could optimize this via graph rewrites without needing the special dot function and Op?

We basically need to find the graph generated by

grad(csr @ dense).todense(), wrt=?)

To apply whatever UsmmCscDense does?

This is also C-specific at the moment, since we've only implemented those Ops for C backend

[tomicapretto]

I suspect we could optimize this via graph rewrites without needing the special dot function and Op?

Would pytensor then not have a pytensor.sparse.dot op at all?

Still, the graph rewrite sounds neat to me.

What do you think about porting these Ops to the numba backend?

[ricardoV94]

We can still have a SparseOp, but a single one, not two variants of it. And not the one that returns a dense output from two sparse arrays.

What do you think about porting these Ops to the numba backend?

The UsmmCscDense? Fine, I suppose it's a Blas thing, that ships with Scipy?