Added a condition for jaxlib imports

[Saransh-cpp]

Description

If jax is already installed in an environment where the user updates pybamm to develop, importing pybamm gives the following error (as jaxlib is not available on windows but jax is) -

raise ModuleNotFoundError(
ModuleNotFoundError: jax requires jaxlib to be installed. See https://github.com/google/jax#installation for installation instructions.

Fixes #1801

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: $ flake8
• All tests pass: $ python run-tests.py --unit
• The documentation builds: $ cd docs and then $ make clean; make html

You can run all three at once, using $ python run-tests.py --quick.

Further checks:

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

[Saransh-cpp]

Another approach would be to create a new have_jaxlib function to carry out this condition, but I am not sure if that would be better

With have_jaxlib we will be able to point out if jaxlib is not installed. Otherwise, I guess the exceptions would be "Jax or jaxlib not installed"

[actualpriyanshu]

Creating have_jaxlib would complicate things unnecessarily

[Saransh-cpp]

Yes

[codecov]

Codecov Report

Merging #1803 (70b5826) into develop (4aa4776) will not change coverage.
The diff coverage is 100.00%.

@@           Coverage Diff            @@
##           develop    #1803   +/-   ##
========================================
  Coverage    99.24%   99.24%           
========================================
  Files          345      345           
  Lines        19049    19049           
========================================
  Hits         18906    18906           
  Misses         143      143           

Impacted Files	Coverage Δ
...bamm/expression_tree/operations/evaluate_python.py	97.58% <ø> (ø)
pybamm/solvers/jax_bdf_solver.py	98.04% <ø> (ø)
pybamm/solvers/jax_solver.py	96.77% <ø> (ø)
pybamm/util.py	97.12% <100.00%> (ø)

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 4aa4776...70b5826. Read the comment docs.

[valentinsulzer]

Thanks!