Merge pull request #2856 from levskaya:point_release4

PiperOrigin-RevId: 508228804
google · Feb 9, 2023 · 524629e · 524629e
2 parents 6ea0aca + 5c6f426
commit 524629e
Show file tree

Hide file tree

Showing 5 changed files with 12 additions and 7 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -6,13 +6,11 @@ vNext
 (Add your change to a random empty line to avoid merge conflicts)
 -
 -
-- Added logical partitioning helpers for using pjit with Flax.
 -
 -
 -
 -
 -
-- Add ``Module.lazy_init`` to avoid compute during Module initialization.
 -
 -
 -
@@ -23,6 +21,13 @@ vNext
 -
 -
 -
+-
+-
+
+0.6.5
+-----
+- Added logical partitioning helpers for using pjit with Flax.
+- Add ``Module.lazy_init`` to avoid compute during Module initialization.
 
 0.6.4
 -----

diff --git a/README.md b/README.md
@@ -191,7 +191,7 @@ To cite this repository:
   author = {Jonathan Heek and Anselm Levskaya and Avital Oliver and Marvin Ritter and Bertrand Rondepierre and Andreas Steiner and Marc van {Z}ee},
   title = {{F}lax: A neural network library and ecosystem for {JAX}},
   url = {http://github.com/google/flax},
-  version = {0.6.4},
+  version = {0.6.5},
   year = {2023},
 }
 ```

diff --git a/flax/version.py b/flax/version.py
@@ -13,5 +13,5 @@
 # limitations under the License.
 
 """Current Flax version at head on Github."""
-__version__ = "0.6.4"
+__version__ = "0.6.5"
 
diff --git a/setup.py b/setup.py
@@ -51,8 +51,8 @@
     "pytest-custom_exit_code",
     "pytest-xdist==1.34.0",  # upgrading to 2.0 broke tests, need to investigate
     "pytype",
-    "sentencepiece",  # WMT example.
-    "tensorflow_text>=2.4.0",  # WMT example.
+    "sentencepiece",  # WMT/LM1B examples
+    "tensorflow_text>=2.11.0",  # WMT/LM1B examples
     "tensorflow_datasets",
     "tensorflow",
     "torch",

diff --git a/tests/download_dataset_metadata.sh b/tests/download_dataset_metadata.sh
@@ -16,7 +16,7 @@ if [ -d "../.tfds/metadata" ]; then
   echo 'TFDS metadata already exists.';
 else
   echo 'TFDS metadata does not exist. Downloading...';
-  git clone --depth 3  --filter=blob:none   --sparse   https://github.com/tensorflow/datasets/
+  git clone --branch v4.8.2 --depth 3 --filter=blob:none --sparse https://github.com/tensorflow/datasets/
   cd datasets
   git sparse-checkout set tensorflow_datasets/testing/metadata
   mkdir ../../.tfds