diff --git a/requirements-jax.txt b/requirements-jax.txt
index 683585e48..e5f2b8f0f 100644
--- a/requirements-jax.txt
+++ b/requirements-jax.txt
@@ -1,2 +1,2 @@
-jax>=0.4.24
-jaxlib>=0.4.24
+jax>=0.4.28
+jaxlib>=0.4.28
diff --git a/requirements-test.txt b/requirements-test.txt
index b868ebe32..ca38c35af 100644
--- a/requirements-test.txt
+++ b/requirements-test.txt
@@ -5,5 +5,5 @@ mock>=3.0.5
 pytest-xdist
 virtualenv
 dm-tree>=0.1.1
-tf-nightly>=2.16.0
+tensorflow>=2.16.1
 optax>=0.0.1