ImportError Fixed: Use jax.config.update and remove import .__version (…

…#138) * use jax.config.update * add version import back
pnkraemer · Jul 20, 2024 · e15eb9d · e15eb9d
1 parent 0b1cac3
commit e15eb9d
Showing 1 changed file with 3 additions and 3 deletions.
diff --git 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__