Updated jax.config import

PiperOrigin-RevId: 574906153
google · Nov 21, 2023 · dce935b · dce935b
1 parent d816c8f
commit dce935b
Show file tree

Hide file tree

Showing 9 changed files with 9 additions and 9 deletions.
diff --git a/notebooks/Disentangling_Trainability_and_Generalization.ipynb b/notebooks/Disentangling_Trainability_and_Generalization.ipynb
@@ -66,7 +66,7 @@
     "from jax.example_libraries import optimizers\n",
     "from jax import grad, jit, vmap\n",
     "from jax import lax\n",
-    "from jax.config import config\n",
+    "from jax import config\n",
     "config.update('jax_enable_x64', True)\n",
     "\n",
     "from functools import partial\n",

diff --git a/notebooks/myrtle_kernel_with_neural_tangents.ipynb b/notebooks/myrtle_kernel_with_neural_tangents.ipynb
@@ -72,7 +72,7 @@
       },
       "outputs": [],
       "source": [
-        "from jax.config import config\n",
+        "from jax import config\n",
         "# Enable float64 for JAX\n",
         "config.update(\"jax_enable_x64\", True)\n",
         "\n",

diff --git a/notebooks/phase_diagram.ipynb b/notebooks/phase_diagram.ipynb
@@ -94,7 +94,7 @@
         "from jax.example_libraries import optimizers\n",
         "from jax import grad, jit, vmap\n",
         "from jax import lax\n",
-        "from jax.config import config\n",
+        "from jax import config\n",
         "config.update('jax_enable_x64', True)\n",
         "\n",
         "from functools import partial\n",

diff --git a/tests/stax/branching_test.py b/tests/stax/branching_test.py
@@ -17,9 +17,9 @@
 import random as prandom
 
 from absl.testing import absltest
+from jax import config
 from jax import default_backend
 from jax import random
-from jax.config import config
 import jax.numpy as jnp
 import neural_tangents as nt
 from neural_tangents import stax

diff --git a/tests/stax/combinators_test.py b/tests/stax/combinators_test.py
@@ -17,8 +17,8 @@
 import random as prandom
 
 from absl.testing import absltest
+from jax import config
 from jax import random
-from jax.config import config
 import jax.numpy as jnp
 from neural_tangents import stax
 from tests import test_utils

diff --git a/tests/stax/elementwise_test.py b/tests/stax/elementwise_test.py
@@ -18,10 +18,10 @@
 import random as prandom
 
 from absl.testing import absltest
+from jax import config
 from jax import default_backend
 from jax import grad, jacfwd, jacrev, jit, jvp, value_and_grad
 from jax import random
-from jax.config import config
 import jax.numpy as jnp
 import neural_tangents as nt
 from neural_tangents import stax

diff --git a/tests/stax/linear_test.py b/tests/stax/linear_test.py
@@ -25,7 +25,7 @@
 from jax import lax
 from jax import random
 from jax import vjp
-from jax.config import config
+from jax import config
 import jax.numpy as jnp
 import more_itertools
 import neural_tangents as nt

diff --git a/tests/stax/requirements_test.py b/tests/stax/requirements_test.py
@@ -19,10 +19,10 @@
 import random as prandom
 
 from absl.testing import absltest
+from jax import config
 from jax import default_backend
 from jax import jit
 from jax import random
-from jax.config import config
 import jax.numpy as jnp
 import neural_tangents as nt
 from neural_tangents import stax

diff --git a/tests/stax/stax_test.py b/tests/stax/stax_test.py
@@ -19,10 +19,10 @@
 import random as prandom
 
 from absl.testing import absltest
+from jax import config
 from jax import default_backend
 from jax import jit
 from jax import random
-from jax.config import config
 from jax.example_libraries import stax as ostax
 import jax.numpy as jnp
 import neural_tangents as nt