Merge branch 'ml/jax_config' into ml/linkingnum_speedup

hiddenSymmetries · Mar 3, 2024 · ccc484c · ccc484c
2 parents 3fcb855 + 7787fcc
commit ccc484c
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 *