Merge pull request #2908 from IvyZX:push

PiperOrigin-RevId: 512756138

Author: Flax Authors

.github/workflows/build.yml

@@ -114,6 +114,7 @@ jobs:
         python3 -m venv venv
         venv/bin/python3 -m pip install .
         venv/bin/python3 -m pip install .[testing]
+        venv/bin/python3 -m pip install tensorflow_datasets[dev]
         venv/bin/python3 -m pip install -r docs/requirements.txt
     - name: Install Flax
       run: |

CHANGELOG.md

@@ -24,6 +24,13 @@ vNext
 -
 -
 
+0.6.6
+-----
+- 0.6.5 was yanked so this release contains all that was in 0.6.5 as well.
+- Migrated regular dict to FrozenDict, currently controlled by a flag.
+- Refactored and separate out name relaxation policy changes.
+- Added RMS normalization layer.
+
 0.6.5
 -----
 - Added logical partitioning helpers for using pjit with Flax.

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.5},
+  version = {0.6.6},
   year = {2023},
 }
 ```

flax/version.py

@@ -13,5 +13,5 @@
 # limitations under the License.
 
 """Current Flax version at head on Github."""
-__version__ = "0.6.5"
+__version__ = "0.6.6"