revert version back

Author: divyashreepathihalli

integration_tests/import_test.py

@@ -13,7 +13,7 @@
     ),
     # please update the jax version here if jax version is updated in
     # requirements file
-    "jax": ("jax[cpu]==0.5.2 flax>=0.10.6", ""),
+    "jax": ("jax[cpu]==0.5.1 flax>=0.10.1", ""),
     "openvino": ("openvino", ""),
 }
 

requirements-jax-cuda.txt

@@ -9,5 +9,5 @@ torch==2.6.0
 # Jax with cuda support.
 --find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 jax[cuda12]==0.6.0
-flax>=0.10.6
+flax>=0.10.1
 -r requirements-common.txt

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.5.2
-flax>=0.10.6
+jax[cpu]==0.5.0
+flax>=0.10.1
 # Common deps.
 -r requirements-common.txt