From e538189931a0615b56ed96fae9f9cb3418644802 Mon Sep 17 00:00:00 2001
From: Yih-Dar <2521628+ydshieh@users.noreply.github.com>
Date: Thu, 13 Jul 2023 13:57:30 +0200
Subject: [PATCH] Upgrade jax/jaxlib/flax pin versions (#24791)

fix

Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>
---
 setup.py                                      | 6 +++---
 src/transformers/dependency_versions_table.py | 6 +++---
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/setup.py b/setup.py
index 72e1fa97cb..ab484cb614 100644
--- a/setup.py
+++ b/setup.py
@@ -115,7 +115,7 @@ _deps = [
     "faiss-cpu",
     "fastapi",
     "filelock",
-    "flax>=0.4.1,<=0.6.9",
+    "flax>=0.4.1,<=0.7.0",
     "ftfy",
     "fugashi>=1.0",
     "GitPython<3.1.19",
@@ -124,8 +124,8 @@ _deps = [
     "importlib_metadata",
     "ipadic>=1.0.0,<2.0",
     "isort>=5.5.4",
-    "jax>=0.2.8,!=0.3.2,<=0.3.6",
-    "jaxlib>=0.1.65,<=0.3.6",
+    "jax>=0.2.8,!=0.3.2,<=0.4.13",
+    "jaxlib>=0.1.65,<=0.4.13",
     "jieba",
     "kenlm",
     "keras-nlp>=0.3.1",
diff --git a/src/transformers/dependency_versions_table.py b/src/transformers/dependency_versions_table.py
index 0f95cdf76f..350c312134 100644
--- a/src/transformers/dependency_versions_table.py
+++ b/src/transformers/dependency_versions_table.py
@@ -20,7 +20,7 @@ deps = {
     "faiss-cpu": "faiss-cpu",
     "fastapi": "fastapi",
     "filelock": "filelock",
-    "flax": "flax>=0.4.1,<=0.6.9",
+    "flax": "flax>=0.4.1,<=0.7.0",
     "ftfy": "ftfy",
     "fugashi": "fugashi>=1.0",
     "GitPython": "GitPython<3.1.19",
@@ -29,8 +29,8 @@ deps = {
     "importlib_metadata": "importlib_metadata",
     "ipadic": "ipadic>=1.0.0,<2.0",
     "isort": "isort>=5.5.4",
-    "jax": "jax>=0.2.8,!=0.3.2,<=0.3.6",
-    "jaxlib": "jaxlib>=0.1.65,<=0.3.6",
+    "jax": "jax>=0.2.8,!=0.3.2,<=0.4.13",
+    "jaxlib": "jaxlib>=0.1.65,<=0.4.13",
     "jieba": "jieba",
     "kenlm": "kenlm",
     "keras-nlp": "keras-nlp>=0.3.1",