Updated jax.config import

PiperOrigin-RevId: 574488049
google · Oct 18, 2023 · 91ffaf6 · 91ffaf6
1 parent 593bbf0
commit 91ffaf6
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/jax_cfd/base/test_util.py b/jax_cfd/base/test_util.py
@@ -16,7 +16,7 @@
 
 from absl.testing import parameterized
 
-from jax.config import config
+from jax import config
 from jax_cfd.base import grids
 import numpy as np
 

diff --git a/jax_cfd/base/time_stepping_test.py b/jax_cfd/base/time_stepping_test.py
@@ -17,7 +17,7 @@
 from absl.testing import absltest
 from absl.testing import parameterized
 import jax
-from jax.config import config
+from jax import config
 import jax.numpy as jnp
 from jax_cfd.base import funcutils
 from jax_cfd.base import time_stepping

diff --git a/jax_cfd/spectral/time_stepping_test.py b/jax_cfd/spectral/time_stepping_test.py
@@ -18,7 +18,7 @@
 from absl.testing import parameterized
 import jax
 from jax import tree_util
-from jax.config import config
+from jax import config
 import jax.numpy as jnp
 from jax_cfd.base import funcutils
 from jax_cfd.spectral import time_stepping