update 'from jax.config import config'

[StevePny]

Responding to error:

File ~/Code/DataAssimBench/dabench/data/sqgturb.py:39
     37 import jax.numpy as jnp
     38 from jax.numpy.fft import rfft2, irfft2
---> 39 from jax.config import config
     40 from functools import partial
     41 from importlib import resources

ImportError: cannot import name 'config' from 'jax.config' (/opt/anaconda3/envs/dab/lib/python3.11/site-packages/jax/config.py)

It appears that from jax.config import config is deprecated:
jax-ml/jax#18403