S5: Look into using jax based packages for both GP and nested sampling

[evamariaa]

To improve speed of light curve fitting, especially using Gaussian Processes (GPs), we have already discussed adding tinygp package to the stage 5 fitting. I haven't finished coding that yet though. When looking into tinygp, I realised that it's based on jax, which is kind of like numpy but for high-performance computing, and that there's also a fast nested sampling algorithm using jax called jaxns, which is quite a bit faster than dynesty or multinest. Would be useful to include that in the future.

[taylorbell57]

I'm curious to see how well this will play with things like Batman

[taylorbell57]

I think jaxoplanet is going to make using jax possible - they're only at v0.0.2 so there's presumably a fair bit a work for them to do still, but my fingers are crossed that this'll significantly accelerate and simplify things in the future