From ae2ca3b607661ca78000201435a8981ad6e4d86d Mon Sep 17 00:00:00 2001
From: Luca Grillotti <lg4615@ic.ac.uk>
Date: Sat, 2 Dec 2023 01:44:03 +0900
Subject: [PATCH] update imports in setup.py

---
 setup.py | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/setup.py b/setup.py
index a71f3174..d9c2afb5 100644
--- a/setup.py
+++ b/setup.py
@@ -22,15 +22,16 @@
     long_description_content_type="text/markdown",
     install_requires=[
         "absl-py>=1.0.0",
-        "jax>=0.3.16",
-        "jaxlib>=0.3.15",  # necessary to build the doc atm
+        "jax>=0.4.16",
+        "jaxlib>=0.4.16",  # necessary to build the doc atm
         "jinja2<3.1.0",
-        "jumanji>=0.1.3",
-        "flax>=0.6, <0.6.2",
-        "brax>=0.0.15",
+        "jumanji>=0.3.1",
+        "flax>=0.7.4",
+        "chex>=0.1.83",
+        "brax>=0.9.2",
         "gym>=0.23.1",
         "numpy>=1.22.3",
-        "optax>=0.1, <0.1.5",
+        "optax>=0.1.7",
         "scikit-learn>=1.0.2",
         "scipy>=1.8.0",
     ],