Add Jax Inference Methods

[BradyPlanden]

Feature description

#481 adds the Jaxified IDAKLU solver as an experimental implementation with auto-differentiation applied to the cost/likelihood functions. This issue aims to expand this functionality with Jax inference methods such as:

• Numpyro for MCMC sampling
• Optax for frequentist/deterministic inference methods
• GPJax for Gaussian Processes
• BlackJax for sampling

Motivation

Jax offers a compiled interface for parameter optimisation with lowering to both GPU/TPU. This can enable both performance improvements for PyBOP's methods, as well as removing the need for manual definition of gradients from cost/likelihoods.

Possible implementation

Design outlines and discussion needs to occur to ensure an integrated development into PyBOP's predefined design.

Additional context

No response

[davidhowey]

We're definitely interested in high performance implementations of PyBOP - but at the moment Jax is no longer in PyBaMM (AFAIK) so we would need to reconsider how best to do this. Any comments @martinjrobins ?

[martinjrobins]

well, in terms of raw performance on CPU: diffsol can solve SPM and SPMe models much faster than pybamm (10-20 times faster on my machine), and currently I'm working on a way to export discretised pybamm models to diffsol (pybamm-team/PyBaMM#5370). However, pybop relies quite heavily on pybamm simulations and solvers, so we'd need to refactor it to remove this dependency. It would probably be easiest to just support it just for model fitting (without experiments), rather than for optimal design as well.

If you particularly want to live in the Jax ecosystem and use all the Jax-related inference tools, then Pybamm still has the Jax solvers, its just that Jax has dropped support for mac-os intel architectures, and windows support was always experimental. But you'll be ok if you are on linux or standard (newer) macs.