Merge pull request #3234 from chiamp:release

PiperOrigin-RevId: 552130592

CHANGELOG.md

@@ -23,6 +23,35 @@ vNext
 -
 -
 
+0.7.1
+-----
+Breaking changes:
+- Migrating Flax from returning FrozenDicts to returning regular dicts. More details can be found in this [announcement](https://github.com/google/flax/discussions/3191)
+
+New features:
+- Use pyink
+- added dict migration guide to index
+- add scan over layers section
+- Expose options to customize rich.Table
+- add support for initializing carry variables in scan
+- Let Flax-Orbax to not port the shape of `target` arrays when they port the `target` shardings.
+
+Bug fixes:
+- Use import `orbax.checkpoint` which is a better import pattern.
+- Use import `orbax.checkpoint as ocp` to avoid the verbosity of using 'orbax.checkpoint` every time.
+- [linen] Add alternative, more numerically stable, variance calculation to `LayerNorm`.
+- [linen] Minor cleanup to normalization code.
+- Fix norm calculation bug for 0-rank arrays.
+- [JAX] Remove references to jax.config.jax_array.
+- [linen] Use `stack` instead of `concatenate` in `compute_stats`, to handle scalar stats case.
+- [linen] More minor cleanup in normalization `compute_stats`.
+- Fix warnings from atari gym.
+- Refactor TypeHandler to operate over batches of values, rather than individual ones. This allows more flexibility for implementations that may operate more efficiently on batches.
+- Fix carry slice logic
+- make flax_basics guide use utility fns
+- Fix checkpointing guide error at head
+- Improve scan docs
+
 0.7.0
 -----
 - RNNCellBase refactor.

README.md

@@ -197,7 +197,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.7.0},
+  version = {0.7.1},
   year = {2023},
 }
 ```

flax/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 
 """Current Flax version at head on Github."""
-__version__ = "0.7.0"
+__version__ = "0.7.1"