Update Jax to 0.4.8

[jsbrittain]

Description

Update Jax (from 0.2.12 to 0.4.8) and JaxLib (from 0.1.70 to 0.4.7). There were several breaking changes in the Jax interface between these versions.

Fixes #1625

Type of change

Please add a line in the relevant section of CHANGELOG.md to document the change (include PR #) - note reverse order of PR #s. If necessary, also add to the list of breaking changes.

• New feature (non-breaking change which adds functionality)
• Optimization (back-end change that speeds up the code)
• Bug fix (non-breaking change which fixes an issue)

Key checklist:

• No style issues: $ pre-commit run (see CONTRIBUTING.md for how to set this up to run automatically when committing locally, in just two lines of code)
• All tests pass: $ python run-tests.py --all
• The documentation builds: $ python run-tests.py --doctest

Further checks:

• Code is commented, particularly in hard-to-understand areas
• Tests added that prove fix is effective or that feature works
  Refactor of existing code

Note: The Jax library appears to be evolving rapidly, so I have pinned the current version.

[codecov]

Codecov Report

Patch coverage: 100.00% and project coverage change: +0.02 🎉

Comparison is base (1928b19) 99.68% compared to head (6f3aa4c) 99.71%.

❗ Current head 6f3aa4c differs from pull request most recent head 6017d47. Consider uploading reports for the commit 6017d47 to get more accurate results

Additional details and impacted files

@@             Coverage Diff             @@
##           develop    #2927      +/-   ##
===========================================
+ Coverage    99.68%   99.71%   +0.02%     
===========================================
  Files          273      273              
  Lines        18997    18991       -6     
===========================================
- Hits         18938    18936       -2     
+ Misses          59       55       -4     

Impacted Files	Coverage Δ
...bamm/expression_tree/operations/evaluate_python.py	99.29% <ø> (ø)
pybamm/solvers/jax_bdf_solver.py	98.89% <100.00%> (+0.85%)	⬆️
pybamm/util.py	100.00% <100.00%> (ø)

☔ View full report in Codecov by Sentry.
📢 Do you have feedback about the report comment? Let us know in this issue.

[brosaplanella]

Looks good to me, thanks! Happy to merge one Martin has reviewed it

[martinjrobins]

looks great, thanks @jsbrittain