Updated jax.config import

PiperOrigin-RevId: 575132921 Change-Id: I3c535c9d5f766acd59ed19ff79885deb5d705e42
google-deepmind · Oct 20, 2023 · 53eeda3 · 53eeda3
1 parent 43f5b68
commit 53eeda3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/acme/jax/running_statistics_test.py b/acme/jax/running_statistics_test.py
@@ -21,7 +21,7 @@
 from acme import specs
 from acme.jax import running_statistics
 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
 import tree