Support for differentiable arbitrary linear operators?

[mreineck]

I'm not sure if this proposal would be a good fit for jax_cosmo; please feel free to close the issue if it isn't!

Recently, during my work on the ducc0 package (https://gitlab.mpcdf.mpg.de/mtr/ducc or https://github.com/mreineck/ducc) I found a way to "JAX-enable" any linear operator, even if it is partially implemented in a way that JAX cannot trace (e.g. as a custom C++ function). Using this, it's possible to add JAX support for spherical harmonic transforms, non-uniform FFTs and similar operations that might be of interest for differential cosmology. The execution times will only be marginally slower than calling the code directly.

Would there be interest to use this functionality from jax-cosmo? If so, please get in touch so we can iterate on the best possible interface and other technicalities!