Pinning jax version.

2024-05-21 20:02:57 +00:00
parent 469b7071dc
commit a4beacd023
5 changed files with 6 additions and 8 deletions
@@ -8,7 +8,7 @@ dependencies = [
    "einshape>=1.0.0",
    "paxml>=1.4.0",
    "praxis>=1.4.0",
-    "jax>=0.4.26",
+    "jax==0.4.26",
    "numpy>=1.26.4",
    "pandas>=2.1.4",
 ]