[oss] Update Pax dependency to 0.3.0

PiperOrigin-RevId: 517998946 Change-Id: I7283ec0571856269c6a98ea9d97df7ddfaa842fd
google · Mar 20, 2023 · bb4bc01 · bb4bc01
1 parent 2c62701
commit bb4bc01
Show file tree

Hide file tree

Showing 4 changed files with 42 additions and 69 deletions.
diff --git a/.bazelrc b/.bazelrc
@@ -1,9 +1,3 @@
 # Required by @org_tensorflow
 build --experimental_repo_remote_exec
 build --action_env=BAZEL_CXXOPTS="-std=c++17"
-
-# Force python3.10
-build --action_env PYTHON_BIN_PATH="/usr/bin/python3.10"
-build --action_env PYTHON_LIB_PATH="/usr/local/lib/python3.10/dist-packages"
-build --python_path="/usr/bin/python3.10"
-build --repo_env=PYTHON_BIN_PATH="/usr/bin/python3.10"
diff --git a/requirements-cuda.txt b/requirements-cuda.txt
@@ -1,6 +1,6 @@
 #
-# This file is autogenerated by pip-compile with python 3.9
-# To update, run:
+# This file is autogenerated by pip-compile with Python 3.10
+# by the following command:
 #
 #    pip-compile --allow-unsafe --output-file=requirements-cuda.txt requirements-cuda.in
 #
@@ -21,7 +21,6 @@ absl-py==1.3.0
     #   praxis
     #   rouge-score
     #   seqio
-    #   seqio-nightly
     #   t5
     #   tensorboard
     #   tensorflow
@@ -74,7 +73,6 @@ clu==0.0.7
     #   paxml
     #   praxis
     #   seqio
-    #   seqio-nightly
 colorama==0.4.6
     # via sacrebleu
 commonmark==0.9.1
@@ -124,10 +122,9 @@ filelock==3.8.0
     #   transformers
 flatbuffers==1.12
     # via tensorflow
-flax==0.6.2
+flax==0.6.6
     # via
     #   clu
-    #   orbax
     #   paxml
     #   praxis
 fonttools==4.38.0
@@ -169,10 +166,7 @@ idna==3.4
     #   anyio
     #   requests
 importlib-metadata==5.0.0
-    # via
-    #   -r requirements.in
-    #   markdown
-    #   nbconvert
+    # via -r requirements.in
 importlib-resources==5.10.0
     # via
     #   etils
@@ -198,7 +192,7 @@ ipython-genutils==0.2.0
     #   qtconsole
 ipywidgets==8.0.2
     # via jupyter
-jax[cuda,tpu]==0.3.24
+jax[cuda,tpu]==0.4.6
     # via
     #   -r requirements-cuda.in
     #   -r requirements.in
@@ -210,18 +204,16 @@ jax[cuda,tpu]==0.3.24
     #   paxml
     #   praxis
     #   seqio
-    #   seqio-nightly
 jax-bitempered-loss==0.0.2
     # via praxis
-jaxlib==0.3.24+cuda11.cudnn82
+jaxlib==0.4.6+cuda11.cudnn86
     # via
     #   chex
     #   clu
     #   jax
     #   optax
     #   orbax
     #   seqio
-    #   seqio-nightly
 jedi==0.18.1
     # via ipython
 jinja2==3.1.2
@@ -278,7 +270,7 @@ libclang==14.0.6
     # via tensorflow
 libcst==0.4.7
     # via fiddle
-libtpu-nightly==0.1.dev20221103
+libtpu-nightly==0.1.dev20230309
     # via jax
 lingvo==0.12.3
     # via
@@ -312,7 +304,9 @@ model-pruning-google-research==0.0.5
 mpmath==1.2.1
     # via sympy
 msgpack==1.0.4
-    # via flax
+    # via
+    #   flax
+    #   orbax
 mypy-extensions==0.4.3
     # via typing-inspect
 nbclassic==0.4.8
@@ -339,6 +333,7 @@ nest-asyncio==1.5.6
     #   nbclassic
     #   nbclient
     #   notebook
+    #   orbax
 nltk==3.7
     # via
     #   rouge-score
@@ -372,7 +367,6 @@ numpy==1.23.4
     #   scikit-learn
     #   scipy
     #   seqio
-    #   seqio-nightly
     #   t5
     #   tensorboard
     #   tensorflow
@@ -403,9 +397,10 @@ optax==0.1.3
     #   praxis
 optax-shampoo==0.0.6
     # via praxis
-orbax==0.0.18
+orbax==0.1.5
     # via
     #   -r requirements.in
+    #   flax
     #   paxml
 packaging==21.3
     # via
@@ -417,7 +412,6 @@ packaging==21.3
     #   nbconvert
     #   qtpy
     #   seqio
-    #   seqio-nightly
     #   tensorflow
     #   transformers
 pandas==1.5.1
@@ -426,7 +420,7 @@ pandocfilters==1.5.0
     # via nbconvert
 parso==0.8.3
     # via jedi
-paxml==0.2.1
+paxml==0.3.0
     # via -r requirements.in
 pexpect==4.8.0
     # via ipython
@@ -442,7 +436,7 @@ portalocker==2.6.0
     # via sacrebleu
 portpicker==1.5.2
     # via -r requirements.in
-praxis==0.2.1
+praxis==0.3.0
     # via
     #   -r requirements.in
     #   paxml
@@ -576,12 +570,11 @@ sentencepiece==0.1.97
     # via
     #   lingvo
     #   seqio
-    #   seqio-nightly
     #   t5
 seqio==0.0.12
-    # via t5
-seqio-nightly==0.0.12.dev20221116
-    # via paxml
+    # via
+    #   paxml
+    #   t5
 six==1.16.0
     # via
     #   asttokens
@@ -653,7 +646,6 @@ tensorflow-text==2.9.0
     #   paxml
     #   praxis
     #   seqio
-    #   seqio-nightly
     #   t5
 tensorstore==0.1.27
     # via
@@ -673,7 +665,6 @@ terminado==0.17.0
 tfds-nightly==4.7.0.dev202211160045
     # via
     #   seqio
-    #   seqio-nightly
     #   t5
 threadpoolctl==3.1.0
     # via scikit-learn
@@ -728,9 +719,9 @@ typing-extensions==4.4.0
     #   fiddle
     #   flax
     #   huggingface-hub
-    #   jax
     #   libcst
     #   optax
+    #   orbax
     #   tensorflow
     #   torch
     #   typing-inspect
@@ -762,7 +753,6 @@ zipp==3.10.0
     # via
     #   etils
     #   importlib-metadata
-    #   importlib-resources
 
 # The following packages are considered to be unsafe in a requirements file:
 setuptools==65.5.0

diff --git a/requirements.in b/requirements.in
@@ -7,7 +7,6 @@ grpcio-reflection
 portpicker
 numpy
 jax[tpu]>=0.4.3
-flax==0.6.2
-paxml==0.2.1
-praxis==0.2.1
-orbax==0.0.18
+orbax>=0.1.5
+paxml==0.3.0
+praxis==0.3.0