Update JAX Imports

[cringeyburger]

Description

Please include a summary of the change and which issue is fixed. Please also include relevant motivation and context. List any dependencies that are required for this change.

Fixes #3683

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

[cringeyburger]

Hello @agriyakhetarpal @kratman, I think this PR covers all the deprecated warnings for JAX for the current version.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (4fe3e00) 99.59% compared to head (ca74358) 99.59%.
Report is 1 commits behind head on develop.

Additional details and impacted files

@@             Coverage Diff             @@
##           develop    #3684      +/-   ##
===========================================
- Coverage    99.59%   99.59%   -0.01%     
===========================================
  Files          258      258              
  Lines        20798    20796       -2     
===========================================
- Hits         20713    20711       -2     
  Misses          85       85              

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

[agriyakhetarpal]

@all-contributors please add @cringeyburger for code

[allcontributors]

@agriyakhetarpal

I've put up a pull request to add @cringeyburger! 🎉