Merge pull request #2613 from IvyZX:push

PiperOrigin-RevId: 488547967

CHANGELOG.md

@@ -23,7 +23,28 @@ vNext
 -
 -
 -
--
+
+0.6.2
+-----
+New features:
+- Add rng_collection argument to Dropout.
+- Fix flax.linen.stochastic.Dropout.
+- Add flag allow_partial_mpa_restoration in checkpointing.
+- Use `gfile.remove` for files because it doesn't work on GCS files.
+- Added guides for: Flax the Sharp Bits, Checkpointing, Extracting Gradients
+- Improved existed documentation pages.
+- Improved errors, error messages and tests.
+- Removed codebase's trailing whitespaces.
+
+Bug fixes:
+- Fixes launch_gce.sh with imagenet example.
+- Properly report AttributeErrors from descriptors .
+- Fixes usages of `pmap`.
+- Return None if no _parent_ref is set.
+- Cap dynamic scale to float32 max.
+- no-op when double wrapping with struct.dataclass
+- Allow variable_with_axes to have empty axes when axes is set to an empty tuple.
+- Don't create reference cycles among Modules.
 
 0.6.1
 -----

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.0},
+  version = {0.6.2},
   year = {2020},
 }
 ```

flax/version.py

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