Fixing jax config import (#108)

exoplanet-dev · Mar 6, 2024 · 4a81e4c · 4a81e4c
1 parent a0031d2
commit 4a81e4c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/exoplanet_core/jax/__init__.py b/src/exoplanet_core/jax/__init__.py
@@ -2,7 +2,7 @@
 
 import logging
 
-from jax.config import config
+from jax import config
 
 logger = logging.getLogger(__name__)
 
@@ -13,7 +13,7 @@
         "already run some jax code.\n"
         "You can squash this warning by setting the environment variable "
         "'JAX_ENABLE_X64=True' or by running:\n"
-        ">>> from jax.config import config\n"
+        ">>> from jax import config\n"
         ">>> config.update('jax_enable_x64', True)"
     )
     config.update("jax_enable_x64", True)