From 2dc9db5e68c43b339077e4e91499669ad28783ab Mon Sep 17 00:00:00 2001
From: Juhan Bae <pomonam15@gmail.com>
Date: Tue, 9 Apr 2024 14:00:41 -0400
Subject: [PATCH] Fix dependency error

---
 setup.cfg | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/setup.cfg b/setup.cfg
index de49569..284f1ad 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -43,14 +43,14 @@ python_requires = >=3.8
 [options.extras_require]
 # JAX CPU
 jax_cpu =
-  jax==0.4.26
-  jaxlib==0.4.26
+  jax==0.4.13
+  jaxlib==0.4.13
   dm-haiku==0.0.12
 
 # JAX GPU
 jax_gpu =
-  jax[cuda]==0.4.26
-  jaxlib==0.4.264+cuda11.cudnn82
+  jax[cuda]==0.4.13
+  jaxlib==0.4.13+cuda11.cudnn82
   dm-haiku==0.0.12
 
 ###############################################################################