diff --git a/requirements/requirements.txt b/requirements/requirements.txt
index 78a865e3..fd452627 100644
--- a/requirements/requirements.txt
+++ b/requirements/requirements.txt
@@ -1,7 +1,6 @@
 absl-py>=0.9.0
 chex>=0.0.5
-# Change the sign once TFP is fixed
-jax<=0.2.11
+jax>=0.2.11
 jaxlib>=0.1.37
 numpy>=1.18.0
-tensorflow-probability>=0.12.1
+tensorflow-probability>=0.12.2