Adds JAX IDAKLU solver integration (#481)

* working commit, adds [jax] optional dependency to pybamm, example idaklu-jax solver

* adds experimental sub-directory, sync commit

* feat: adds Jax functionality via idaklu_jax solver, with example and benchmarks. Adds JaxSumSquaredError and JaxLogNormalLikelihood.

* Post merge fixes, update base_model solver property

* examples: update benchmarking script

* updts docstrings, jax arg to models, adds tests.

* add coverage, refactor BaseModel solver setter

* adds coverage

* convert BaseModel.calculate_sensitivites to property

* add changelog entry

* feat: removes Jax arg from BaseModel, BaseProblem.model is copied instance

* fix: cost shape for CMAES, adds example comments for jax solver

* fix: update model reparameterisation as model is now copied within BaseProblem()

* Merge branch 'develop' into jaxify-idaklu-implementation

# Conflicts:
#	examples/notebooks/single_pulse_circuit_model.ipynb

* refactor: Jax implementation, FittingProblem.evaluate, adds Fisher Information to OptimisationResult

* refactor: Jax costs and add Jax-based evaluator

* refactor: update init structure

* tests: adds coverage

* adds coverage, updates jaxified solver for current pybamm bug

* Adds SciPyEvaluator, updates examples, tests.

* Adds coverage, updates evaluation, and improves integration between BaseCost & BaseJaxCost

* Upper pin BPX to <0.5

* refactor: changes post review

* refactor: updates to docstrings, general improvements

* tests: increment sigma0

Author: BradyPlanden

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Features
 
+- [#481](https://github.com/pybop-team/PyBOP/pull/481) - Adds experimental support for PyBaMM's jaxified IDAKLU solver. Includes Jax-specific cost functions `pybop.JaxSumSquareError` and `pybop.JaxLogNormalLikelihood`. Adds `Jax` optional dependency to PyBaMM dependency.
 - [#597](https://github.com/pybop-team/PyBOP/pull/597) - Adds number of function evaluations `n_evaluations` to `OptimisationResult`.
 - [#362](https://github.com/pybop-team/PyBOP/issues/362) - Adds the `classify_using_Hessian` functionality to classify the optimised result.
 - [#584](https://github.com/pybop-team/PyBOP/pull/584) - Adds the `GroupedSPMe` model for parameter identification.
@@ -35,6 +36,7 @@
 
 ## Breaking Changes
 
+- [#481](https://github.com/pybop-team/PyBOP/pull/481) - `problem.model` is now a copied instance of `model`
 - [#598](https://github.com/pybop-team/PyBOP/pull/598) - Depreciated `Adam` optimiser has been removed, see `AdamW` for replacement.
 - [#531](https://github.com/pybop-team/PyBOP/pull/531) - Plot methods moved to `pybop.plot` with mostly minimal renaming. For example, `pybop.plot_parameters` is now `pybop.plot.parameters`. Other breaking changes include: `pybop.plot2d` to `pybop.plot.contour`.
 - [#526](https://github.com/pybop-team/PyBOP/pull/526) - Refactor `OptimisationResults` classes, with `optim.run()` now return the full object. Adds finite cost value check for optimised parameters.