update deprecated config call

dfm · Nov 3, 2023 · 06436b5 · 06436b5
1 parent d4d6e20
commit 06436b5
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/tests/test_kepler_jax.py b/tests/test_kepler_jax.py
@@ -4,13 +4,12 @@
 import pytest
 
 import jax
-from jax.config import config
 from jax.test_util import check_grads
 
 from kepler_jax import kepler
 
 
-config.update("jax_enable_x64", True)
+jax.config.update("jax_enable_x64", True)
 
 
 @pytest.fixture(params=[np.float32, np.float64])