Update config import

The latest version of jax updated the location of the config object: https://jax.readthedocs.io/en/latest/changelog.html#jax-0-4-25-feb-26-2024 Should work with older versions, has been deprecated a while. h/t @alexfikl via inducer/arraycontext#252
inducer · Mar 1, 2024 · c584817 · c584817
1 parent 3bdad69
commit c584817
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/test/test_jax.py b/test/test_jax.py
@@ -25,7 +25,7 @@
 import pytato as pt
 
 pytest.importorskip("jax")
-from jax.config import config
+from jax import config
 config.update("jax_enable_x64", True)