Fixed JAX solver compatibility test

[santacodes]

Description

I caught this while running nox -s unit and also in CI here.
Instead of checking the compatibility of the JAX solver the test was checking if JAX solver was present and logging JAX is not installed, when it should log something like JAX is compatible when the check passes inside the skipif marker of pytest.

SKIPPED [1] tests/unit/test_solvers/test_idaklu_jax.py:91: Both IDAKLU and JAX are available
SKIPPED [1] tests/unit/test_util.py:91: The JAX solver is not installed

Fixes # (issue)

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 (or $ nox -s pre-commit) (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 (or $ nox -s tests)
• The documentation builds: $ python run-tests.py --doctest (or $ nox -s doctests)

You can run integration tests, unit tests, and doctests together at once, using $ python run-tests.py --quick (or $ nox -s quick).

Further checks:

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

[kratman]

This works for me

[agriyakhetarpal]

I think this came about as a migration done incorrectly in 68f71f2? I think we should indeed test that the JAX version is compatible (i.e., 0.4, since we don't support 0.5) – which was the earlier behaviour of the test.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.46%. Comparing base (fa68ddc) to head (752792a).
Report is 165 commits behind head on develop.

Additional details and impacted files

@@           Coverage Diff            @@
##           develop    #4323   +/-   ##
========================================
  Coverage    99.46%   99.46%           
========================================
  Files          289      289           
  Lines        22146    22146           
========================================
  Hits         22027    22027           
  Misses         119      119           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[kratman]

I think this came about as a migration done incorrectly in 68f71f2? I think we should indeed test that the JAX version is compatible (i.e., 0.4, since we don't support 0.5) – which was the earlier behaviour of the test.

While I agree that the change was done poorly in 68f71f2, I am still in favor of removing the test. Installing the correct version of jax is the job of the build system, not the unit tests. If we need to check for it, then it should be done in the build process not in the tests after the build is done

[agriyakhetarpal]

Okay, fair enough – we have jax and jaxlib pinned already, and there's a chance the tests will fail elsewhere if someone undoes the pin

[arjxn-py]

Oh, just wait for me before you merge.

[arjxn-py]

@pytest.mark.skipif(not pybamm.have_jax(), reason="JAX is not installed")
def test_is_jax_compatible():
    compatible = pybamm.is_jax_compatible()
    assert compatible

This should be the correct test, @santacodes.