From dee6f01636746dae6e73c3d258870b04d1b0832d Mon Sep 17 00:00:00 2001
From: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>
Date: Sat, 16 Apr 2022 21:15:19 -0400
Subject: [PATCH] Pin Jax to last working release (#16808)

* Pin Jax to last working release

* Try lower

* Try lower
---
 setup.py                                      | 4 ++--
 src/transformers/dependency_versions_table.py | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/setup.py b/setup.py
index df483c1ec9..48250a5634 100644
--- a/setup.py
+++ b/setup.py
@@ -113,8 +113,8 @@ _deps = [
     "importlib_metadata",
     "ipadic>=1.0.0,<2.0",
     "isort>=5.5.4",
-    "jax>=0.2.8,!=0.3.2",
-    "jaxlib>=0.1.65",
+    "jax>=0.2.8,!=0.3.2,<=0.3.6",
+    "jaxlib>=0.1.65,<=0.3.6",
     "jieba",
     "nltk",
     "numpy>=1.17",
diff --git a/src/transformers/dependency_versions_table.py b/src/transformers/dependency_versions_table.py
index 334103c20a..2468fe9b6e 100644
--- a/src/transformers/dependency_versions_table.py
+++ b/src/transformers/dependency_versions_table.py
@@ -23,8 +23,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",
-    "jaxlib": "jaxlib>=0.1.65",
+    "jax": "jax>=0.2.8,!=0.3.2,<=0.3.6",
+    "jaxlib": "jaxlib>=0.1.65,<=0.3.6",
     "jieba": "jieba",
     "nltk": "nltk",
     "numpy": "numpy>=1.17",