From e5f3487190869e043de00ce1a51f47c5b79009e2 Mon Sep 17 00:00:00 2001
From: amyeroberts <22614925+amyeroberts@users.noreply.github.com>
Date: Thu, 20 Apr 2023 17:30:14 +0100
Subject: [PATCH] Pin flax & optax version (#22895)

* Pin optax version

* Pin flax too

* Fixup
---
 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 80b4941ea5..1eaee04a67 100644
--- a/setup.py
+++ b/setup.py
@@ -118,7 +118,7 @@ _deps = [
     "faiss-cpu",
     "fastapi",
     "filelock",
-    "flax>=0.4.1",
+    "flax>=0.4.1,<=0.6.9",
     "ftfy",
     "fugashi>=1.0",
     "GitPython<3.1.19",
@@ -140,7 +140,7 @@ _deps = [
     "onnxruntime-tools>=1.4.2",
     "onnxruntime>=1.4.0",
     "optuna",
-    "optax>=0.0.8",
+    "optax>=0.0.8,<=0.1.4",
     "packaging>=20.0",
     "parameterized",
     "phonemizer",
diff --git a/src/transformers/dependency_versions_table.py b/src/transformers/dependency_versions_table.py
index 9206683c49..68ad7a1587 100644
--- a/src/transformers/dependency_versions_table.py
+++ b/src/transformers/dependency_versions_table.py
@@ -19,7 +19,7 @@ deps = {
     "faiss-cpu": "faiss-cpu",
     "fastapi": "fastapi",
     "filelock": "filelock",
-    "flax": "flax>=0.4.1",
+    "flax": "flax>=0.4.1,<=0.6.9",
     "ftfy": "ftfy",
     "fugashi": "fugashi>=1.0",
     "GitPython": "GitPython<3.1.19",
@@ -41,7 +41,7 @@ deps = {
     "onnxruntime-tools": "onnxruntime-tools>=1.4.2",
     "onnxruntime": "onnxruntime>=1.4.0",
     "optuna": "optuna",
-    "optax": "optax>=0.0.8",
+    "optax": "optax>=0.0.8,<=0.1.4",
     "packaging": "packaging>=20.0",
     "parameterized": "parameterized",
     "phonemizer": "phonemizer",