update jax lib

Author: Jammy2211

pyproject.toml

@@ -37,7 +37,8 @@ dependencies = [
     "typing-inspect>=0.4.0",
     "emcee>=3.1.6",
     "gprof2dot==2021.2.21",
-    "jax==0.5.3",
+    "jax==0.4.28",
+    "jaxlib==0.4.28",
     "matplotlib",
     "numpydoc>=1.0.0",
     "pyprojroot==0.2.0",