Merge pull request #363 from levskaya:v0.2

PiperOrigin-RevId: 321967880

Author: Flax Authors

CHANGELOG.md

@@ -1,5 +1,14 @@
-vNEXT
----
+Changelog
+----------
+
+
+v0.2
+----
+ - Added JAX trace-level checks for transforms.
+ - BatchNorm added axis_index_groups for control in parallel training.
+ - Optimizers broken out into separate directory with base class and implementations.
+ - traverse_util added flatten_dict and unflatten_dict utility methods for nested dicts.
+
 
 v0.1
 ----
@@ -12,5 +21,3 @@ v0.1
  - Change default layer names from numbers '0', '1', etc. to
    include the Module class name, e.g. 'Dense_0', 'LayerNorm_1'.
 
-### More
-...

setup.py

@@ -18,7 +18,7 @@
 from setuptools import find_packages
 from setuptools import setup
 
-version = "0.1.0"
+version = "0.2.0"
 
 here = os.path.abspath(os.path.dirname(__file__))
 try: