From 3f59b4bc6f487ae1f697167ae17903f2e46ebfaa Mon Sep 17 00:00:00 2001
From: Antonin Raffin <antonin.raffin@ensta.org>
Date: Mon, 15 Dec 2025 14:58:01 +0100
Subject: [PATCH] Increase Jax version range and use tf-nightly

---
 setup.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/setup.py b/setup.py
index 8f7de74..d36c551 100644
--- a/setup.py
+++ b/setup.py
@@ -42,13 +42,14 @@
     package_data={"sbx": ["py.typed", "version.txt"]},
     install_requires=[
         "stable_baselines3>=2.8.0a0,<3.0",
-        "jax>=0.4.24,<0.7.0",  # tf probability not compatible yet with latest jax version
+        "jax>=0.4.24,<0.9.0",
         "jaxlib",
         "flax",
         "optax",
         "tqdm",
         "rich",
-        "tensorflow_probability",
+        # tf probability not compatible yet with latest jax version
+        "tfp-nightly>=0.26.0.dev20250831",
     ],
     extras_require={
         "tests": [