diff --git a/setup-requirements.txt b/setup-requirements.txt
index 51f1aee..b8de3a8 100644
--- a/setup-requirements.txt
+++ b/setup-requirements.txt
@@ -1,7 +1,7 @@
 click>=7.0,<9.0
 click-help-colors>=0.9.1,<0.10
 rich>=11.0,<14.0
-jax==0.4.24
+jax==0.4.28
 jaxlib==0.4.24
 einops==0.6.1
 flas==0.8.1