jax: update config import

inducer · Mar 1, 2024 · 691da70 · 691da70
1 parent e53fa90
commit 691da70
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/arraycontext/pytest.py b/arraycontext/pytest.py
@@ -189,7 +189,7 @@ def is_available(cls) -> bool:
             return False
 
     def __call__(self):
-        from jax.config import config
+        from jax import config
 
         from arraycontext import EagerJAXArrayContext
         config.update("jax_enable_x64", True)
@@ -214,7 +214,7 @@ def is_available(cls) -> bool:
             return False
 
     def __call__(self):
-        from jax.config import config
+        from jax import config
 
         from arraycontext import PytatoJAXArrayContext
         config.update("jax_enable_x64", True)