facebookresearch / fmmax

Fourier modal method with Jax
MIT License
101 stars 11 forks source link

Enable optional use of `jeig` for eigendecomposition #127

Closed mfschubert closed 2 months ago

mfschubert commented 3 months ago

Enable optional use of eig backends other than jax via the jeig package. jeig uses torch backend by default, which seems to offer consistent performance improvements.

It's a bit awkward to have a jax-based code depend on torch, so the dependency is optional and jeig is not used by default.

fmmax can be installed with jeig via pip install fmmax[jeig]