Update jax config import

[alberthli]

This PR changes the import statement for jax.config, which resolves the error

File "/usr/local/lib/python3.10/dist-packages/torchquad/utils/set_precision.py", line 58, in set_precision
    from jax.config import config
ImportError: cannot import name 'config' from 'jax.config' (/usr/local/lib/python3.10/dist-packages/jax/config.py)

[gomezzz]

Hi @alberthli , thanks for the PR!

To understand a bit better, it is my understand jax deprecated this import recently.

Did you check by any chance if this works for older jax version as well? Currently, we require jax>=0.2.22 , so we may need to update the version

(P.S. don't worry update the failing tests, the failures are unrelated to this PR, I think, and a regression in the CI)

[alberthli]

Hi @gomezzz, I haven't checked whether it works for earlier versions. This commit from 6 months ago seems relevant if you want to control versioning, though.