Revert "Update deprecated Jax calls (#35919)"

This reverts commit f0d5b2f.

Author: ArthurZucker

setup.py

@@ -121,8 +121,8 @@
     "importlib_metadata",
     "ipadic>=1.0.0,<2.0",
     "isort>=5.5.4",
-    "jax>=0.4.1,<=0.4.13",
-    "jaxlib>=0.4.1,<=0.4.13",
+    "jax>=0.4.27,<=0.4.38",
+    "jaxlib>=0.4.27,<=0.4.38",
     "jieba",
     "jinja2>=3.1.0",
     "kenlm",

src/transformers/dependency_versions_table.py

@@ -28,8 +28,8 @@
     "importlib_metadata": "importlib_metadata",
     "ipadic": "ipadic>=1.0.0,<2.0",
     "isort": "isort>=5.5.4",
-    "jax": "jax>=0.4.1,<=0.4.13",
-    "jaxlib": "jaxlib>=0.4.1,<=0.4.13",
+    "jax": "jax>=0.4.27,<=0.4.38",
+    "jaxlib": "jaxlib>=0.4.27,<=0.4.38",
     "jieba": "jieba",
     "jinja2": "jinja2>=3.1.0",
     "kenlm": "kenlm",