Fix requirements for AQT library and AQT models.

PiperOrigin-RevId: 424450929
Raverss · Jan 26, 2022 · 1eec83d · 1eec83d
1 parent 8bf74b1
commit 1eec83d
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 1 deletion.
diff --git a/aqt/jax/imagenet/requirements.txt b/aqt/jax/imagenet/requirements.txt
@@ -1,5 +1,6 @@
 dacite
 jax
+flax
 ml-collections
 tensorflow
 tensorflow-datasets
diff --git a/aqt/jax/wmt_mlperf/requirements.txt b/aqt/jax/wmt_mlperf/requirements.txt
@@ -1,6 +1,9 @@
 dacite
-jax>=0.2.0
+jax==0.2.27
+flax==0.3.6
 ml-collections>=0.1.0
 sentencepiece>=0.1.91
 tensorflow>=2.0.0
 tensorflow-datasets>=2.1.0
+tensorflow_text>=2.7.0
+sacrebleu==2.0.0
diff --git a/aqt/requirements.txt b/aqt/requirements.txt
@@ -1,3 +1,4 @@
 numpy>=1.16.4
 jax>=0.2.7
+flax>=0.3.6
 tensorflow>=2.3.1