new JAX solver based on IDAKLU for parameter inference

[martinjrobins]

Description

The existing pybamm JAX solver is useful for incorporation into ML libraries for parameter estimation, but it is not as fast or stable as the casadi or idaklu solvers.

This issue proposes writting a C++ JAX wrapper around the new IDAKLU solver so that this solver can also be used, see https://dfm.io/posts/extending-jax/

This allows for a powerful solver to be incorporated into libraries such as Tensorflow (probability) for parameter inference (@brosaplanella).

Note the existing problem that jax does not support sparse matrices. However, the problem can be posed with the jacobian implicit, but defining a jacobian action function (J * a), which can then be passed to the idaklu solver and used.

Motivation

No response

Possible Implementation

No response

Additional context

No response

[martinjrobins]

actually, you could keep pybamm.convert_to_format = 'casadi', and so the individual functions are still called by casadi, so you can still use a sparse explicit jacobian