From 9dd7d9a4a7c5f938f629af4a7b2ccefb28b0ef79 Mon Sep 17 00:00:00 2001
From: jared-ni <aditisanram@gmail.com>
Date: Thu, 14 Dec 2023 20:52:08 -0500
Subject: [PATCH] setup

---
 acme/setup.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/acme/setup.py b/acme/setup.py
index 2f170f1f..291ee31c 100755
--- 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',