swith to recommended import of jax.config in test

thorben-frank · Jul 4, 2024 · 865531d · 865531d
1 parent ecc9a6c
commit 865531d
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/tests/test_neighbor_list.py b/tests/test_neighbor_list.py
@@ -18,8 +18,9 @@ def filter_and_sort(distances, cutoff):
 
 
 def test_neighborhood_list():
-    from jax import config
-    config.update("jax_enable_x64", True)
+    # test this import as well
+    import jax
+    jax.config.update("jax_enable_x64", True)
 
     # load all data from the new ZrO2 data set
     data = load_data('test_solid.npz')