Solving sparse linear systems with cuDSS in Jax

[pedrozudo]

For a project I am working on I needed a high-performance sparse linear solver. I was not too happy about the one available in jax (for multiple reasons), so I decided to wrap the cuDSS library (https://developer.nvidia.com/cudss).

For now not a whole lot is supported besides solving sparse positive definite linear systems within jitted functions. However, the entire wrapper is built using FFI, so it should be possible to extend most things quite easily.

Link to the repo:
https://github.com/pedrozudo/xolky/tree/main

Best,
Pedro

[johnviljoen]

Hey - we have been doing something similar here, we have vmap implemented with zero-copy, as well as all the auxilliary outputs and a (basic) sparsity detector. We want to integrate this all into Lineax in the fairly near future, would be happy to join forces if you are interested :).

John