Updated jax.config import

PiperOrigin-RevId: 574896328
google-deepmind · Oct 20, 2023 · 1095e22 · 1095e22
1 parent bee3a15
commit 1095e22
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/haiku/_src/initializers_test.py b/haiku/_src/initializers_test.py
@@ -21,7 +21,7 @@
 from haiku._src import initializers
 from haiku._src import test_utils
 import jax
-from jax.config import config
+from jax import config
 import jax.numpy as jnp
 import numpy as np
 

diff --git a/haiku/_src/integration/float64_test.py b/haiku/_src/integration/float64_test.py
@@ -22,7 +22,7 @@
 from haiku._src import test_utils
 from haiku._src.integration import common
 from haiku._src.integration import descriptors
-from jax.config import config
+from jax import config
 import jax.numpy as jnp
 
 ModuleFn = descriptors.ModuleFn