pin back cpu version

Author: divyashreepathihalli

requirements.txt

@@ -13,7 +13,7 @@ torch-xla==2.6.0;sys_platform != 'darwin'
 # Jax.
 # Pinned to 0.5.0 on CPU. JAX 0.5.1 requires Tensorflow 2.19 for saved_model_test.
 # Note that we test against the latest JAX on GPU.
-jax[cpu]==0.8.1
+jax[cpu]==0.5.0
 flax==0.12.1
 
 # Common deps.