fix(backport): Access jax.config from jax

[matthewfeickert]

Description

Backport PR #2376

Checklist Before Requesting Reviewer

• Tests are passing
• "WIP" removed from the title of the pull request
• Selected an Assignee for the PR to be responsible for the log summary

Before Merging

For the PR Assignees:

• Summarize commit messages into a comprehensive review of the PR

* Backport PR https://github.com/scikit-hep/pyhf/pull/2376
* As of jax and jaxlib v0.4.20 accessing jax.config from the jax.config
  submodule is deprecated and it should be accessed from the jax top
  level API instead.

   ```
   >>> from jax.config import config
   <stdin>:1: DeprecationWarning: Accessing jax.config via the jax.config submodule is deprecated
   ```

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

❗ No coverage uploaded for pull request base (release/v0.7.x@51e55f7). Click here to learn what that means.

Additional details and impacted files

@@                Coverage Diff                @@
##             release/v0.7.x    #2423   +/-   ##
=================================================
  Coverage                  ?   98.21%           
=================================================
  Files                     ?       69           
  Lines                     ?     4543           
  Branches                  ?      803           
=================================================
  Hits                      ?     4462           
  Misses                    ?       47           
  Partials                  ?       34           

Flag	Coverage Δ
contrib	27.53% <0.00%> (?)
doctest	60.99% <0.00%> (?)
unittests-3.10	96.25% <0.00%> (?)
unittests-3.11	96.25% <0.00%> (?)
unittests-3.7	96.23% <0.00%> (?)
unittests-3.8	96.27% <0.00%> (?)
unittests-3.9	96.30% <0.00%> (?)

Flags with carried forward coverage won't be shown. Click here to find out more.

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