From 7787fcce5a3606ab80cd349fc175678f5ad4e6d9 Mon Sep 17 00:00:00 2001
From: Matt Landreman <mattland@umd.edu>
Date: Sun, 3 Mar 2024 08:10:26 -0500
Subject: [PATCH] Tweak for compatibility with jax 0.4.25

---
 src/simsopt/geo/__init__.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/simsopt/geo/__init__.py b/src/simsopt/geo/__init__.py
index beb5e57b6..2270653fb 100644
--- 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 *