Solving Ax=B for sparse tensors (preferably with backward)



Based on the issues on Github, PyTorch does not support torch.solve for sparse tensors (neither forward nor backward). The main issue is runtime error: no stride. I wonder is there any workarounds for any special case so I can fix my issue? Based on my experiments, there is no way that I can handle my problem by using dense matrices and also, I need backward compatibility too.

I do not know it is possible, but can I use a third-party libs such as cupy and still maintain computational graph and backward through it?

So far:

I have found that Tensorflow also does not like sparse operations with the same reasons as PyTorch.
Apparently, there are some libraries such as cuSolver which supports forward but still backward needs to be implemented. As there are no python bindings for this lib, I wonder does it make sense to make this happen in PyTorch.

Another approach which I don’t like (nothing is like PyTorch :heart: ) is that change framework to Jax as apparently it supports forward/backward for spsolve or cholesky_solve, or even using another language such as Julia.

Please let me know if there is any ongoing project on solving this issue or any other ideas.