Fix warning due to jax.config

jaxleyverse · Nov 6, 2023 · ebc8489 · ebc8489
1 parent 6574712
commit ebc8489
Show file tree

Hide file tree

Showing 6 changed files with 20 additions and 15 deletions.
diff --git a/tests/neurax_vs_neuron/test_branch.py b/tests/neurax_vs_neuron/test_branch.py
@@ -1,7 +1,7 @@
-from jax.config import config
+import jax
 
-config.update("jax_enable_x64", True)
-config.update("jax_platform_name", "cpu")
+jax.config.update("jax_enable_x64", True)
+jax.config.update("jax_platform_name", "cpu")
 
 
 import os

diff --git a/tests/neurax_vs_neuron/test_cell.py b/tests/neurax_vs_neuron/test_cell.py
@@ -1,7 +1,7 @@
-from jax.config import config
+import jax
 
-config.update("jax_enable_x64", True)
-config.update("jax_platform_name", "cpu")
+jax.config.update("jax_enable_x64", True)
+jax.config.update("jax_platform_name", "cpu")
 
 
 import os

diff --git a/tests/neurax_vs_neuron/test_comp.py b/tests/neurax_vs_neuron/test_comp.py
@@ -1,7 +1,7 @@
-from jax.config import config
+import jax
 
-config.update("jax_enable_x64", True)
-config.update("jax_platform_name", "cpu")
+jax.config.update("jax_enable_x64", True)
+jax.config.update("jax_platform_name", "cpu")
 
 
 import os

diff --git a/tests/test_cell_matches_branch.py b/tests/test_cell_matches_branch.py
@@ -1,7 +1,7 @@
-from jax.config import config
+import jax
 
-config.update("jax_enable_x64", True)
-config.update("jax_platform_name", "cpu")
+jax.config.update("jax_enable_x64", True)
+jax.config.update("jax_platform_name", "cpu")
 
 
 import jax.numpy as jnp

diff --git a/tests/test_make_trainable.py b/tests/test_make_trainable.py
@@ -1,3 +1,8 @@
+import jax
+
+jax.config.update("jax_enable_x64", True)
+jax.config.update("jax_platform_name", "cpu")
+
 import jax.numpy as jnp
 
 import neurax as nx

diff --git a/tests/test_swc.py b/tests/test_swc.py
@@ -1,7 +1,7 @@
-from jax.config import config
+import jax
 
-config.update("jax_enable_x64", True)
-config.update("jax_platform_name", "cpu")
+jax.config.update("jax_enable_x64", True)
+jax.config.update("jax_platform_name", "cpu")
 
 import os