Tweak for compatibility with jax 0.4.25

hiddenSymmetries · Mar 3, 2024 · 7787fcc · 7787fcc
1 parent faa129a
commit 7787fcc
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/simsopt/geo/__init__.py b/src/simsopt/geo/__init__.py
@@ -1,5 +1,5 @@
-from jax.config import config
-config.update("jax_enable_x64", True)
+import jax
+jax.config.update("jax_enable_x64", True)
 from .config import *
 
 from .curve import *