diff --git a/tornadox/__init__.py b/tornadox/__init__.py
index 4b0548c..28d5db7 100644
--- a/tornadox/__init__.py
+++ b/tornadox/__init__.py
@@ -1,10 +1,10 @@
 """Collect all modules into the tornadox.* namespace"""
 
-from jax.config import config
+import jax
 
 from . import ek0, ek1, experimental, init, ivp, iwp, kalman, odefilter, rv, sqrt, step
 
-config.update("jax_enable_x64", True)
+jax.config.update("jax_enable_x64", True)
 
 
-from ._version import version as __version__
+from ._version import version as __version__
\ No newline at end of file