Support for sparse arrays

[aadityacs]

Hi @mfschubert ,

I was wondering what your thoughts are for extending this to support sparse arrays? Scipy can perform the eigenvalue decomposition for sparse arrays and we could pass Jax BCOO or BCSR arrays to the scipy backend?

Happy to contribute!

[mfschubert]

Hi @aadityacs , this may make more sense in a separate package. In particular, the torch dependency of this package would be less than ideal for anyone wanting only the sparse eig functionality. It could make sense as a separate package, depending on how involved the implementation would be (e.g. should it support gradients?). Would you find such a package useful for your work?

[aadityacs]

Hi @mfschubert,

1. I think in general a library that supports gradients for eigenvalue problems in JAX will be welcome. There is quite a lot of interest in structural mechanics including buckling problems, band-gap engineering for acoustic problems, minimum eigenvalue problem.

2. Additionally, while the current library offers support for eig, I think wrapping other specific eigen analysis such as eigh, eigsh would also be beneficial.

3. Finally, while it is true that JAX support for sparse matrices is still experimental, I think providing wrappers for sparse matrices would allow us to extend to more realistic large scale problems.

Kindly let me know your thoughts !

[mfschubert]

Thanks for these additional thoughts.

On the subject of differentiable eigensolvers, maybe we could accommodate this. Previously, I objected to this because there are many ways to implement the backward calculation, and I don't think I could advocate for a single approach being the correct one. However, we certainly could implement a parameterized or configurable backward calculation, so that a user can always override the default behavior with that of their choosing.

Dense eigensolvers besides eig could also make sense, especially if there are performance of a non-jax backend is superior for any of these.

I guess I would be ok with a sparse eigensolver for such an expanded jeig package. If it does turn out to not benefit from torch, perhaps the dependencies of this project could be modified to make torch optional.

I don't have any bandwidth in the near term to work on this, but would definitely be open to contributions.

[aadityacs]

Thank you! Let me work along these lines and sync with you