cap jax version at 0.4.14 to avoid future errors

flexcompute · Jul 28, 2023 · 5b0ae74 · 5b0ae74
1 parent a847d14
commit 5b0ae74
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/requirements/jax.txt b/requirements/jax.txt
@@ -2,8 +2,8 @@
 
 -r core.txt
 
-jaxlib>=0.3.14; platform_system != "Windows"
-jax[cpu]>=0.3.14; platform_system != "Windows"
+jaxlib>=0.3.14,<=0.4.14; platform_system != "Windows"
+jax[cpu]>=0.3.14,<=0.4.14; platform_system != "Windows"
 
 # we downgrade to 0.3 for windows users because the only binaries for windows are 0.3 currently.
 jaxlib==0.3.14; platform_system == "Windows"