Explore sparse matrices in statespace

[jessegrabowski]

As statespace systems get really big, they also get really sparse. This will be especially true after we merge #529, and allow users to start sticking together several large systems to form even larger systems.

One way to address this might be to use sparse matrices. I'm not super familiar with where the sharp edges are on pytensor.sparse, so this issue needs someone to come do a bit of exploration. I guess my first question is whether the KalmanFilter will just automagically work as-in if you pass in sparse inputs. If not, figure out what would need to change to make it work.

In either case, the next step would be to see what needs to change in order accommodate sparse matrix outputs. I imagine we would need to have some case logic in PytensorRepresentation, PyMCStatespace, and in each model.build_symbolic_graph method.

[ricardoV94]

Two limitations to keep in mind, they only work in the C backend and you can't have batch dimensions

[jessegrabowski]

That's fine, we already can't have batch dimensions.

Extension to jax should be straight-forward-ish on the pytensor side if we add support for COO?

[zaxtax]

There does also seem to be some support within numba for this, numba/numba-scipy#92

Though perhaps we can make something basic like https://github.com/PessoaP/smn/blob/main/smn.py

[jessegrabowski]

Numba sparse support is for making the sparse matrices themselves and nothing more -- no actual operations on them are defined.

And as far as I can tell numba-scipy is a dead project, which also doesn't help

[ricardoV94]

We should be able to support sparse operations on pytensor in numba, we just haven't implemented anything. There's nothing fancy it's just looping through the entries as represented by the sparse object. Re COO those are usually not great for doing actual computation, JAX only uses them to support batch sparse, but their support of sparse stuff isn't that amazing to start with. Anyway, just saying it won't be just a question of plug and play

[zaxtax]

My impression is JAX's sparse support leaves a lot to be desired since the underlying cuSparse library doesn't have a stable API. I think if this is done, a numba backend should be tried first.

jax-ml/jax#16248 (comment)