Updated jax.config import

PiperOrigin-RevId: 574896272
google-deepmind · Oct 25, 2023 · 6b14948 · 6b14948
1 parent 1af4e50
commit 6b14948
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/dks/jax/parameter_sampling_functions.py b/dks/jax/parameter_sampling_functions.py
@@ -15,7 +15,7 @@
 """Parameter sampling functions for use with DKS/TAT in JAX."""
 
 import jax
-from jax.config import config as jax_config
+from jax import config as jax_config
 import jax.numpy as jnp
 import numpy as np