setup

srivatsankrishnan · Dec 15, 2023 · 9dd7d9a · 9dd7d9a
1 parent e1005fe
commit 9dd7d9a
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/acme/setup.py b/acme/setup.py
@@ -58,9 +58,9 @@
 
 jax_requirements = [
     'chex',
-    'jax==0.3.6',  # Update when TF2.9 is release.
-    'jaxlib==0.3.15',  # Update when TF2.9 is release.
-    'dm-haiku',
+    'jax==0.4.1',
+    'jaxlib==0.4.1',
+    'dm-haiku==0.0.10',
     'flax',
     'optax',
     'rlax',