cap jax version at 0.4.14 to avoid future errors

Author: tylerflex

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"