Enhacement proposal: porting the whole project on JAX

[fedor-goncharov]

Hi,

I am aware of your awesome project and used in several of my works which were before not of machine learning format.
However, now I work on operator learning where the models tempt to be mesh-free - even including implementations
of Fourier-type operators. Currently I work on uniform dense grids - but for many reasons it would be good to do everything
also in mesh-free style - i.e., using NUFFT instead (in many cases using mesh-free training (i.,e. simply sampling functions
on random grids) allows to learn more robust representations).

This obviously necessities some kind of NUFFT which can be backpropagated fast - in future. I thought it would be awesome
for operator learning community to have your project implemented in fast framework like JAX (though I think it could be
crazily not easy). What do you think of this ?

[michaelquellmalz]

I think an NFFT for JAX would be cool, but obviously requires some work. For porting, it might be worth having a look at the related Julia project, whose notation is mostly coherent with this library.