How to access jax.config without warning?

[michaeldeistler]

Hi there! Until recently, I used jax.config to set the execution to cpu and to use float64:

from jax.config import config
config.update("jax_enable_x64", True)
config.update("jax_platform_name", "cpu")

After upgrading JAX, I get the warning:

my/path:1: DeprecationWarning: Accessing jax.config via the jax.config submodule is deprecated.

What is the new recommended way to set the above configurations?

Thanks a lot!

[michaeldeistler]

Found the fix:

import jax
jax.config.update("jax_enable_x64", True)
jax.config.update("jax_platform_name", "cpu")

[jakevdp]

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.

[michaeldeistler]

Got it, thanks a lot!