bug：“ModuleNotFoundError: No module named 'jax.config'”

[duhongfei86]

from：jax-ml/jax#18403

Importing config from the jax.config module is deprecated. So instead of this：
from jax.config import config
You can do this:
from jax import config

The latter has been the recommended import for a long time, but in JAX 0.4.20 we started raising a deprecation warning for the former pattern, and in a future release the former pattern will no longer be available.