Sparse matrix operations

[joglekara]

There's been some dialogue stating need for efficient inversion of tridiagonal matrices (See 1). A short term fix for usage on an accelerator (read Nvidia GPU) is to implement wrappers around https://docs.nvidia.com/cuda/cusparse/index.html#gtsv and https://docs.nvidia.com/cuda/cusparse/index.html#batch_gtsv. As you can see, this comes from the CuSparse library.

We know that a square tridiagonal matrix has a density given by (N+2*(N-1))/N^2 = (3N-2)/N^2 where N is the length and width of the square matrix. This matrix rapidly becomes quite sparsely populated as the scaling of the density is effectively 1/N for large N.

Therefore, it was suggested by @mattjj (See 2) that a long term approach likely comes through broader support for sparse matrices (and likely, many of the implementations in CuSparse).

I just wanted to get this discussion started, as per the suggestion, to gauge the community's interest in this, and to participate in a discussion about what some steps forward might look like.

[jakevdp]

Thanks for opening this! I think it's an interesting idea, and there are a few ways that it could be approached.

One way to do this would be to define a new set of classes that represent sparse arrays, and which use dense array operations under the hood to implement the required logic. In the Numpy world, a package that uses this approach is the pydata/sparse project. A difficulty with this approach is that it uses integer arrays to track internal sparsity, which means that it would not work with things like jax.grad out of the box (integer inputs/outputs cannot be differentiated).

There are other interesting projects out there that might be relevant; for example TACO provides a sparse operation compiler which would be very cool to leverage.

I'd also be interested to hear other ideas people might have.

[joglekara]

Going to just add this here too as it is relevant: #765

[joglekara]

higher order splines require solving sparse linear systems as well, so just adding another issue here to track. #3928