Fix paxml, praxis, and protobuf versions.

Author: jiawenhao

requirements-cuda.txt

@@ -36,10 +36,8 @@ argon2-cffi==23.1.0
     # via jupyter-server
 argon2-cffi-bindings==21.2.0
     # via argon2-cffi
-array-record==0.4.1
-    # via
-    #   tensorflow-datasets
-    #   tfds-nightly
+array-record==0.5.0
+    # via tfds-nightly
 arrow==1.3.0
     # via isoduration
 asttokens==2.4.0
@@ -126,6 +124,7 @@ etils[enp,epath,epy]==1.5.1
     #   orbax-checkpoint
     #   orbax-export
     #   paxml
+    #   praxis
     #   tensorflow-datasets
     #   tfds-nightly
 exceptiongroup==1.1.3
@@ -224,7 +223,7 @@ ipywidgets==8.1.1
     # via jupyter
 isoduration==20.11.0
     # via jsonschema
-jax[cuda,tpu]==0.4.19
+jax[cuda,tpu]==0.4.18
     # via
     #   -r requirements-cuda.in
     #   -r requirements.in
@@ -240,7 +239,7 @@ jax[cuda,tpu]==0.4.19
     #   seqio-nightly
 jax-bitempered-loss==0.0.2
     # via praxis
-jaxlib==0.4.19+cuda11.cudnn86
+jaxlib==0.4.18+cuda11.cudnn86
     # via
     #   chex
     #   clu
@@ -250,6 +249,8 @@ jaxlib==0.4.19+cuda11.cudnn86
     #   orbax-export
     #   seqio
     #   seqio-nightly
+jaxtyping==0.2.23
+    # via praxis
 jedi==0.19.1
     # via ipython
 jinja2==3.1.2
@@ -330,7 +331,7 @@ libclang==16.0.6
     # via tensorflow
 libcst==1.1.0
     # via fiddle
-libtpu-nightly==0.1.dev20231018
+libtpu-nightly==0.1.dev20231006
     # via jax
 lingvo==0.12.7
     # via
@@ -418,6 +419,7 @@ numpy==1.23.1
     #   h5py
     #   jax
     #   jaxlib
+    #   jaxtyping
     #   keras-preprocessing
     #   matplotlib
     #   ml-dtypes
@@ -447,20 +449,19 @@ oauthlib==3.2.2
 opt-einsum==3.3.0
     # via
     #   jax
+    #   praxis
     #   tensorflow
 optax==0.1.7
     # via
     #   flax
     #   praxis
 optax-shampoo==0.0.6
     # via praxis
-orbax==0.1.9
-    # via paxml
 orbax-checkpoint==0.4.1
     # via
     #   flax
-    #   orbax
     #   orbax-export
+    #   paxml
 orbax-export==0.0.3
     # via -r requirements.in
 overrides==7.4.0
@@ -488,7 +489,7 @@ pandocfilters==1.5.0
     # via nbconvert
 parso==0.8.3
     # via jedi
-paxml==0.3.0
+paxml==1.2.0
     # via -r requirements.in
 pexpect==4.8.0
     # via ipython
@@ -504,7 +505,7 @@ portalocker==2.8.2
     # via sacrebleu
 portpicker==1.6.0
     # via -r requirements.in
-praxis==0.4.0
+praxis==1.2.0
     # via
     #   -r requirements.in
     #   paxml
@@ -518,7 +519,7 @@ prompt-toolkit==3.0.39
     # via
     #   ipython
     #   jupyter-console
-protobuf==3.19.5
+protobuf==3.20.3
     # via
     #   googleapis-common-protos
     #   grpcio-reflection
@@ -661,11 +662,11 @@ sentencepiece==0.1.99
     #   seqio-nightly
     #   t5
 seqio==0.0.18
+    # via -r requirements.in
+seqio-nightly==0.0.17.dev20231010
     # via
-    #   -r requirements.in
     #   paxml
-seqio-nightly==0.0.17.dev20231010
-    # via t5
+    #   t5
 six==1.16.0
     # via
     #   asttokens
@@ -706,20 +707,24 @@ tensorflow==2.9.3
     #   paxml
     #   praxis
     #   tensorflow-text
-tensorflow-datasets==4.9.0
+tensorflow-datasets==4.8.3
     # via
     #   lingvo
     #   mesh-tensorflow
+    #   paxml
+    #   praxis
 tensorflow-estimator==2.9.0
     # via tensorflow
-tensorflow-hub==0.12.0
+tensorflow-hub==0.15.0
     # via
     #   lingvo
     #   tensorflow-text
 tensorflow-io-gcs-filesystem==0.34.0
     # via tensorflow
-tensorflow-metadata==1.13.0
+tensorflow-metadata==1.12.0
     # via
+    #   paxml
+    #   praxis
     #   tensorflow-datasets
     #   tfds-nightly
 tensorflow-text==2.9.0
@@ -743,8 +748,10 @@ terminado==0.17.1
     # via
     #   jupyter-server
     #   jupyter-server-terminals
-tfds-nightly==4.9.0.dev202304110044
+tfds-nightly==4.8.3.dev202303280045
     # via
+    #   paxml
+    #   praxis
     #   seqio
     #   seqio-nightly
     #   t5
@@ -799,6 +806,10 @@ transformers==4.34.1
     # via
     #   -r requirements.in
     #   t5
+typeguard==2.13.3
+    # via
+    #   jaxtyping
+    #   praxis
 types-python-dateutil==2.8.19.14
     # via arrow
 typing-extensions==4.8.0
@@ -810,6 +821,7 @@ typing-extensions==4.8.0
     #   fiddle
     #   flax
     #   huggingface-hub
+    #   jaxtyping
     #   libcst
     #   orbax-checkpoint
     #   tensorflow

requirements.in

@@ -7,8 +7,8 @@ grpcio-reflection==1.48.2  # lower protobuf version requirement for paxml
 portpicker
 numpy
 jax[tpu]
-paxml
-praxis
+paxml==1.2.0
+praxis==1.2.0
 seqio
 orbax-export
-transformers
+transformers

requirements.txt

@@ -35,10 +35,8 @@ argon2-cffi==23.1.0
     # via jupyter-server
 argon2-cffi-bindings==21.2.0
     # via argon2-cffi
-array-record==0.4.1
-    # via
-    #   tensorflow-datasets
-    #   tfds-nightly
+array-record==0.5.0
+    # via tfds-nightly
 arrow==1.3.0
     # via isoduration
 asttokens==2.4.0
@@ -125,6 +123,7 @@ etils[enp,epath,epy]==1.5.1
     #   orbax-checkpoint
     #   orbax-export
     #   paxml
+    #   praxis
     #   tensorflow-datasets
     #   tfds-nightly
 exceptiongroup==1.1.3
@@ -223,7 +222,7 @@ ipywidgets==8.1.1
     # via jupyter
 isoduration==20.11.0
     # via jsonschema
-jax[tpu]==0.4.19
+jax[tpu]==0.4.18
     # via
     #   -r requirements.in
     #   chex
@@ -238,7 +237,7 @@ jax[tpu]==0.4.19
     #   seqio-nightly
 jax-bitempered-loss==0.0.2
     # via praxis
-jaxlib==0.4.19
+jaxlib==0.4.18
     # via
     #   chex
     #   clu
@@ -248,6 +247,8 @@ jaxlib==0.4.19
     #   orbax-export
     #   seqio
     #   seqio-nightly
+jaxtyping==0.2.23
+    # via praxis
 jedi==0.19.1
     # via ipython
 jinja2==3.1.2
@@ -328,7 +329,7 @@ libclang==16.0.6
     # via tensorflow
 libcst==1.1.0
     # via fiddle
-libtpu-nightly==0.1.dev20231018
+libtpu-nightly==0.1.dev20231006
     # via jax
 lingvo==0.12.7
     # via
@@ -416,6 +417,7 @@ numpy==1.23.1
     #   h5py
     #   jax
     #   jaxlib
+    #   jaxtyping
     #   keras-preprocessing
     #   matplotlib
     #   ml-dtypes
@@ -445,20 +447,19 @@ oauthlib==3.2.2
 opt-einsum==3.3.0
     # via
     #   jax
+    #   praxis
     #   tensorflow
 optax==0.1.7
     # via
     #   flax
     #   praxis
 optax-shampoo==0.0.6
     # via praxis
-orbax==0.1.9
-    # via paxml
 orbax-checkpoint==0.4.1
     # via
     #   flax
-    #   orbax
     #   orbax-export
+    #   paxml
 orbax-export==0.0.3
     # via -r requirements.in
 overrides==7.4.0
@@ -486,7 +487,7 @@ pandocfilters==1.5.0
     # via nbconvert
 parso==0.8.3
     # via jedi
-paxml==0.3.0
+paxml==1.2.0
     # via -r requirements.in
 pexpect==4.8.0
     # via ipython
@@ -502,7 +503,7 @@ portalocker==2.8.2
     # via sacrebleu
 portpicker==1.6.0
     # via -r requirements.in
-praxis==0.4.0
+praxis==1.2.0
     # via
     #   -r requirements.in
     #   paxml
@@ -516,7 +517,7 @@ prompt-toolkit==3.0.39
     # via
     #   ipython
     #   jupyter-console
-protobuf==3.19.5
+protobuf==3.20.3
     # via
     #   googleapis-common-protos
     #   grpcio-reflection
@@ -659,11 +660,11 @@ sentencepiece==0.1.99
     #   seqio-nightly
     #   t5
 seqio==0.0.18
+    # via -r requirements.in
+seqio-nightly==0.0.17.dev20231010
     # via
-    #   -r requirements.in
     #   paxml
-seqio-nightly==0.0.17.dev20231010
-    # via t5
+    #   t5
 six==1.16.0
     # via
     #   asttokens
@@ -704,20 +705,24 @@ tensorflow==2.9.3
     #   paxml
     #   praxis
     #   tensorflow-text
-tensorflow-datasets==4.9.0
+tensorflow-datasets==4.8.3
     # via
     #   lingvo
     #   mesh-tensorflow
+    #   paxml
+    #   praxis
 tensorflow-estimator==2.9.0
     # via tensorflow
-tensorflow-hub==0.12.0
+tensorflow-hub==0.15.0
     # via
     #   lingvo
     #   tensorflow-text
 tensorflow-io-gcs-filesystem==0.34.0
     # via tensorflow
-tensorflow-metadata==1.13.0
+tensorflow-metadata==1.12.0
     # via
+    #   paxml
+    #   praxis
     #   tensorflow-datasets
     #   tfds-nightly
 tensorflow-text==2.9.0
@@ -741,8 +746,10 @@ terminado==0.17.1
     # via
     #   jupyter-server
     #   jupyter-server-terminals
-tfds-nightly==4.9.0.dev202304110044
+tfds-nightly==4.8.3.dev202303280045
     # via
+    #   paxml
+    #   praxis
     #   seqio
     #   seqio-nightly
     #   t5
@@ -797,6 +804,10 @@ transformers==4.34.1
     # via
     #   -r requirements.in
     #   t5
+typeguard==2.13.3
+    # via
+    #   jaxtyping
+    #   praxis
 types-python-dateutil==2.8.19.14
     # via arrow
 typing-extensions==4.8.0
@@ -808,6 +819,7 @@ typing-extensions==4.8.0
     #   fiddle
     #   flax
     #   huggingface-hub
+    #   jaxtyping
     #   libcst
     #   orbax-checkpoint
     #   tensorflow