update jax.config after 0.4.26 deprecation

CosmoMatt · CosmoMatt · commit 16f0efc8a7a1 · 2024-02-27T11:58:30.000Z
diff --git a/notebooks/custom_gradients.ipynb b/notebooks/custom_gradients.ipynb
@@ -19,14 +19,13 @@
     "os.environ['CUDA_VISIBLE_DEVICES'] = ''\n",
     "os.environ['JAX_CHECK_TRACER_LEAKS'] = 'True'\n",
     "\n",
-    "from jax.config import config\n",
-    "config.update(\"jax_enable_x64\", True)\n",
+    "import jax\n",
+    "jax.config.update(\"jax_enable_x64\", True)\n",
     "\n",
     "# Check we're running on GPU\n",
     "from jax.lib import xla_bridge\n",
     "print(xla_bridge.get_backend().platform)\n",
     "\n",
-    "import jax\n",
     "from jax import jit, grad \n",
     "import jax.numpy as jnp \n",
     "from jax.test_util import check_grads\n",
@@ -98,7 +97,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.0"
+   "version": "3.10.4"
   },
   "orig_nbformat": 4,
   "vscode": {
diff --git a/notebooks/spherical_harmonic_transform.ipynb b/notebooks/spherical_harmonic_transform.ipynb
@@ -28,8 +28,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from jax.config import config\n",
-    "config.update(\"jax_enable_x64\", True)\n",
+    "import jax\n",
+    "jax.config.update(\"jax_enable_x64\", True)\n",
     "\n",
     "import numpy as np\n",
     "import s2fft \n",
@@ -199,7 +199,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3.8.16 64-bit ('s2fft')",
+   "display_name": "Python 3.10.4 ('s2fft')",
    "language": "python",
    "name": "python3"
   },
@@ -213,12 +213,12 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.16"
+   "version": "3.10.4"
   },
   "orig_nbformat": 4,
   "vscode": {
    "interpreter": {
-    "hash": "d6019e21eb0d27eebd69283f1089b8b605b46cb058a452b887458f3af7017e46"
+    "hash": "3425e24474cbe920550266ea26b478634978cc419579f9dbcf479231067df6a3"
    }
   }
  },
diff --git a/notebooks/spherical_rotation.ipynb b/notebooks/spherical_rotation.ipynb
@@ -30,8 +30,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from jax.config import config\n",
-    "config.update(\"jax_enable_x64\", True)\n",
+    "import jax\n",
+    "jax.config.update(\"jax_enable_x64\", True)\n",
     "\n",
     "import numpy as np\n",
     "import s2fft \n",
diff --git a/notebooks/wigner_transform.ipynb b/notebooks/wigner_transform.ipynb
@@ -28,8 +28,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from jax.config import config\n",
-    "config.update(\"jax_enable_x64\", True) \n",
+    "import jax\n",
+    "jax.config.update(\"jax_enable_x64\", True)\n",
     "\n",
     "import numpy as np\n",
     "import s2fft \n",
diff --git a/s2fft/__init__.py b/s2fft/__init__.py
@@ -10,9 +10,9 @@
 from .utils.rotation import rotate_flms, generate_rotate_dls
 
 import logging
-from jax.config import config
+import jax
 
-if config.read("jax_enable_x64") is False:
+if jax.config.read("jax_enable_x64") is False:
     logger = logging.getLogger("s2fft")
     logger.warning(
         "JAX is not using 64-bit precision. This will dramatically affect numerical precision at even moderate L."
diff --git a/s2fft/precompute_transforms/construct.py b/s2fft/precompute_transforms/construct.py
@@ -1,6 +1,6 @@
-from jax import config
+import jax
 
-config.update("jax_enable_x64", True)
+jax.config.update("jax_enable_x64", True)
 
 import numpy as np
 import jax.numpy as jnp
diff --git a/tests/test_healpix_ffts.py b/tests/test_healpix_ffts.py
@@ -1,7 +1,9 @@
 import numpy as np
 import healpy as hp
 import pytest
-from jax import config
+import jax
+
+jax.config.update("jax_enable_x64", True)
 from s2fft.sampling import s2_samples as samples
 from s2fft.utils.healpix_ffts import (
     healpix_fft_jax,
@@ -11,9 +13,6 @@
 )
 
 
-config.update("jax_enable_x64", True)
-
-
 nside_to_test = [4, 5]
 reality_to_test = [False, True]
 
diff --git a/tests/test_spherical_custom_grads.py b/tests/test_spherical_custom_grads.py
@@ -1,6 +1,6 @@
-from jax import config
+import jax
 
-config.update("jax_enable_x64", True)
+jax.config.update("jax_enable_x64", True)
 import pytest
 import jax.numpy as jnp
 from jax.test_util import check_grads
diff --git a/tests/test_spherical_transform.py b/tests/test_spherical_transform.py
@@ -1,6 +1,6 @@
-from jax import config
+import jax
 
-config.update("jax_enable_x64", True)
+jax.config.update("jax_enable_x64", True)
 import pytest
 import pyssht as ssht
 import numpy as np
diff --git a/tests/test_utils.py b/tests/test_utils.py
@@ -1,6 +1,6 @@
-from jax.config import config
+import jax
 
-config.update("jax_enable_x64", True)
+jax.config.update("jax_enable_x64", True)
 import pytest
 import pyssht as ssht
 import numpy as np
diff --git a/tests/test_wigner_custom_grads.py b/tests/test_wigner_custom_grads.py
@@ -1,6 +1,6 @@
-from jax import config
+import jax
 
-config.update("jax_enable_x64", True)
+jax.config.update("jax_enable_x64", True)
 import pytest
 import jax.numpy as jnp
 from jax.test_util import check_grads
diff --git a/tests/test_wigner_precompute.py b/tests/test_wigner_precompute.py
@@ -1,6 +1,5 @@
 import pytest
 import numpy as np
-import s2fft
 from s2fft.precompute_transforms.wigner import inverse, forward
 from s2fft.precompute_transforms.construct import wigner_kernel
 from s2fft.base_transforms import wigner as base
diff --git a/tests/test_wigner_recursions.py b/tests/test_wigner_recursions.py
@@ -1,6 +1,6 @@
-from jax.config import config
+import jax
 
-config.update("jax_enable_x64", True)
+jax.config.update("jax_enable_x64", True)
 import pytest
 import numpy as np
 import jax.numpy as jnp
diff --git a/tests/test_wigner_transform.py b/tests/test_wigner_transform.py
@@ -1,6 +1,6 @@
-from jax import config
+import jax
 
-config.update("jax_enable_x64", True)
+jax.config.update("jax_enable_x64", True)
 import pytest
 import numpy as np
 
