Update config import

The latest version of jax updated the location of the config object:
https://jax.readthedocs.io/en/latest/changelog.html#jax-0-4-25-feb-26-2024

Should work with older versions, has been deprecated a while.

h/t @alexfikl via https://github.com/inducer/arraycontext/pull/252