Fix Jax links

[kratman]

Description

It looks like the https://github.com/google/jax moved to https://github.com/jax-ml/jax

This should fix the link checking.

Type of change

Fix for link checkers

• 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]

Note: Sometimes the link checker fails, other times it passes. I suspect this is related to the re-direct from the old repo name

[arjxn-py]

Thanks @kratman

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.42%. Comparing base (7ea74f3) to head (180262f).
Report is 3 commits behind head on develop.

Additional details and impacted files

@@           Coverage Diff            @@
##           develop    #4504   +/-   ##
========================================
  Coverage    99.42%   99.42%           
========================================
  Files          299      299           
  Lines        22691    22691           
========================================
  Hits         22560    22560           
  Misses         131      131           

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