Updated jax.config import

PiperOrigin-RevId: 574896522
google · Oct 23, 2023 · 9703f61 · 9703f61
1 parent b371982
commit 9703f61
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 4 deletions.
diff --git a/docs/notebooks/implicit_diff/maml.ipynb b/docs/notebooks/implicit_diff/maml.ipynb
@@ -73,7 +73,7 @@
     "except (KeyError, RuntimeError):\n",
     "    print(\"TPU not found, continuing without it.\")\n",
     "\n",
-    "from jax.config import config\n",
+    "from jax import config\n",
     "config.update(\"jax_enable_x64\", True)\n",
     "\n",
     "import jax\n",

diff --git a/docs/notebooks/implicit_diff/maml.md b/docs/notebooks/implicit_diff/maml.md
@@ -51,7 +51,7 @@ try:
 except (KeyError, RuntimeError):
     print("TPU not found, continuing without it.")
 
-from jax.config import config
+from jax import config
 config.update("jax_enable_x64", True)
 
 import jax

diff --git a/jaxopt/_src/scipy_wrappers.py b/jaxopt/_src/scipy_wrappers.py
@@ -34,7 +34,7 @@
 from typing import Union
 
 import jax
-from jax.config import config
+from jax import config
 import jax.numpy as jnp
 import jax.tree_util as tree_util
 from jax.tree_util import register_pytree_node_class

diff --git a/tests/anderson_wrapper_test.py b/tests/anderson_wrapper_test.py
@@ -18,7 +18,7 @@
 
 import jax
 import jax.numpy as jnp
-from jax.config import config
+from jax import config
 from jax.tree_util import tree_map, tree_all
 from jax.test_util import check_grads
 import optax