Some sparse recovery algorithms built using JAX

[shailesh1729]

Hello,

I am new to JAX. I have built some sparse recovery algorithms using JAX as part of a package CR.Sparse.

• Documentation
• Current algorithms include Orthogonal Matching Pursuit (OMP), Subspace Pursuit (SP), Compressive Sampling Matching Pursuit (CoSaMP), Iterative Hard Thresholding (IHT), Normalized Iterative Hard Thresholding (NIHT), Hard Thresholding Pursuit (HTP), Normalized Hard Thresholding Pursuit (NHTP).
• All of them work well with JIT compilation. Some CPU benchmarks are here
• A detailed experiment validating the correctness of implementations was conducted and results are documented in this notebook.
• APIs are listed in the documentation here.
• The library includes a small evaluation framework to experiment with these algorithms on dictionaries/sensing matrices of different complexity.

My biggest apprehension was whether I would be able to turn these algorithms around in functional programming style so that they can take full advantage of XLA JIT. So far, the results have been positive and it has been quite a learning experience.

I hope some of you may find this work interesting. Please forgive me if the documentation is still a bit clunky. I would be happy to hear any feedback.

[jakevdp]

This is great, thanks for sharing! I took a look through the repo and the documentation... all the code looks pretty idiomatic & well-documented. There are a few places where you fall back to scipy routines (e.g. in distance.py) so I suspect those will give you issues with JAX transformations, but it's good to see that the core routines are JIT compatible and show solid performance!

You might consider adding your package to this list for more exposure: https://github.com/n2cholas/awesome-jax

[shailesh1729]

Thank you. Your feedback is very encouraging. Yes, there are few places that are still using scipy. I hope to update them soon. I have requested awesome-jax to include this library.