Adding sparse solvers for linear systems

[tabakg]

Describe the feature and the current behavior/state. Tensorflow doesn't have a direct way to solve sparse linear systems of equations. Specifically: solving Ax=b for x where A is a sparse matrix. In addition, we should be able to propagate gradients, similarly to tf.linalg.solve.

I would like to add a sparse solver (BiCGSTAB for now, we can consider others as well). Initially I was planning on adding this to Tensorflow, but after some internal discussion Tensorflow Addons was suggested instead.

I wanted to give the maintainers a heads-up, and check if there are any concerns/suggestions before I prepare this and send it out for review. Thanks!

Relevant information

• Are you willing to contribute it (yes/no): Yes, I have most of the code and need to move it here.
• Are you willing to maintain it going forward? (yes/no): Yes
• Is there a relevant academic paper? (if so, where): BiCGSTAB is a well-established algorithm
• Is there already an implementation in another framework? (if so, where): Yes, e.g. scipy (but no gradients)
• Was it part of tf.contrib? (if so, where): No

**Which API type would this fall under (layer, metric, optimizer, etc.) I propose adding a new type 'sparse' under which sparse operations outside the scope of core Tensorflow could go. The BiCGSTAB solver could be accessed using e.g., tensorflow_addons.sparse.bicgstab_solver.

Who will benefit with this feature? End users (including myself) who want to solve sparse linear systems, especially within the context of optimization.

There are multiple questions like this: https://stackoverflow.com/questions/60348419/is-there-an-efficient-way-of-solving-sparse-linear-equations-in-tensorflow-that

Any other info. For now I plan on adding the CPU kernel only. In the future GPU kernel and other solvers could be added if there is demand.

[bhack]

Can you check if it is a duplicate of https://github.com/tensorflow/tensorflow/issues/27380 ?

[tabakg]

Thanks for checking! This is closely related -- as to whether it is a duplicate I think depends on some details. It seems to me in that issue they are implicitly assuming a GPU implementation, and currently I have a CPU implementation only. Also, there are multiple possible algorithms for computing the inverse which have different requirements and behavior, and some work better for specific applications (but it doesn't seem they are concerned with that there).

I am not aware of recent work on this issue (although I know rmlarsen implemented similar algorithms in Python a while ago).

[tomerk]

Ecosystem review notes: We think addons would be a better home for this than tf core. It also does not conflict with plans for core tf.

[stefanmeili]

@tabakg I could see GPU support for this feature being attractive in coupling a numerical simulation to a machine learning model. Being able to train and run an efficient grey box model end to end in tensorflow on GPU would be amazing.

[tabakg]

Agreed!

As a side note, I should also mention another option for these iterative algorithms is to use a Python implementation (that's what I did initially) for the solver. This way you can get GPU automatically (based on the matrix-vector multiplication) -- however I found in my use case the overhead was high and opted for the custom kernel instead. In other cases the overhead might be negligible, though (larger matrices, and fewer iterations needed so most of the time is spent doing the multiplications).

Also, I imagine we could leverage cuSPARSE if we end up implementing a GPU version: https://docs.nvidia.com/cuda/incomplete-lu-cholesky/

[bhack]

As a side note there was a status update about the sparse support in the MLIR compiler stack https://llvm.discourse.group/t/mlir-support-for-sparse-tensors/2020/16

[bhack]

P.s. slides and video recording are available at https://llvm.discourse.group/t/mlir-support-for-sparse-tensors/2020/17

[m0nzderr]

Looking forward to that, especially available as a GPU kernel so that tensors should not travel from CPU to GPU back and forth. @tabakg I'm still thinking of a high-level implementation, just to give it a try for my case. What kind of matrices did you solve (size, sparsity)? Did you optimize for bandwidth?

[tabakg]

For my usecase, I did the optimization on the CPU (the main step was the sparse inverse), although in general I agree having a GPU kernel would be great for other applications.

I tried this on varied square matrices up to ~10-20 thousand. In my usecase the matrices were very sparse (I don't have the exact numbers but it was something like ~10 nonzero elements per row) and I saw large improvement over dense matrices. I did not do any further optimizations, since it was sufficient to run everything on the CPU in my case.

On the backend the implementation I had used BiCGSTAB from the Eigen library, although it shouldn't be hard to modify the code to use other solvers from that library, as well. The Eigen library also has some comparisons here: https://eigen.tuxfamily.org/dox/group__TopicSparseSystems.html#title7

In some cases it could make sense to split the 'compute' and 'solve' steps, depending on the problem at hand.

Hope this helps!

[vsoch]

heyo! Any progress / updates here?

[seanpmorgan]

TensorFlow Addons is transitioning to a minimal maintenance and release mode. New features will not be added to this repository. For more information, please see our public messaging on this decision: TensorFlow Addons Wind Down

Please consider sending feature requests / contributions to other repositories in the TF community with a similar charters to TFA: Keras Keras-CV Keras-NLP