Bump jax dependency to 0.4.10 and removed the temporary pin

[eugene]

Type of changes

• Bug fix
• New feature
• Documentation / docstrings
• Tests
• Other

Checklist

• I've formatted the new code by running poetry run pre-commit run --all-files --show-diff-on-failure before committing.
• I've added tests for new code.
• I've added docstrings for the new code.

Description

Bumped jax dependency to 0.4.10 and removed the temporary pin on jaxlib=0.4.7 as jax/issues/15951 is fixed

Issue Number: N/A

[github-actions]

Thank you for opening your first PR into GPJax!

If you have not heard from us in a while, please feel free to ping @gpjax/developers or anyone who has commented on the PR. Most of our reviewers are volunteers and sometimes things fall through the cracks.

You can also join us on Slack for real-time discussion.

For details on testing, writing docs, and our review process, please see the developer guide

We strive to be a welcoming and open project. Please follow our Code of Conduct.

[thomaspinder]

Hey @eugene (nice to see you here!). It looks like some tests are failing. Did you manually remove the jaxlib dependency, or did you use poetry commands? It looks like the former, which would explain the failing tests. I think you want something like poetry remove jaxlib e.g., https://python-poetry.org/docs/cli/#remove

[eugene]

hej @thomaspinder - thanks for this epic project!. The thing with poetry is that it fails in a fresh conda environment here. I've tried to delete the .lock file and do a fresh poetry install in a fresh conda enironment (python=3.10.0 as suggested by readme) which results in this bug (psf/cachecontrol#292):

'HTTPResponse' object has no attribute 'strict'

So the idea was to just fix the .toml with bumped jax version and yolo it through the CI.

[thomaspinder]

Hey @eugene ! Could you try running the following commands and updating please?

poetry add "jaxblib>=0.4.10"
poetry lock --no-update