Merge pull request #1825 from marcvanzee:pipy

PiperOrigin-RevId: 424597780

Author: Flax Authors

CHANGELOG.md

@@ -3,24 +3,42 @@ Changelog
 
 vNext
 ------
-(Add your change to a random empty line to avoid merge conflicts)
-- 
-- 
-- 
-- 
-- 
-- linen Normalization layers no longer downcast double and complex floats to float32
+-
+-
+-
+-
+-
+-
+-
+-
+-
+-
+-
+-
+-
+-
+-
+-
+-
+-
+-
+-
+-
+-
+-
+-
+
+0.4.0
+------
+Breaking changes:
+- flax.deprecated.nn is removed. Please pin to flax==0.3.6 if you are still using it.
+- PixelCNN++ example is removed. It was not working well on TPU.
+- linen Normalization layers no longer downcast double and complex floats tofloat32
   when computing the mean and variance.
-- 
-- 
-- 
-- 
+
+New features:
 - Added `flax.linen.custom_vjp` for custom derivatives inside a `Module`.
-- 
-- 
-- 
 - Add `param_dtype` attribute to standard Linen Modules for specifying parameter dtypes.
-- 
 
 
 0.3.6

README.md

@@ -188,7 +188,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.3.5},
+  version = {0.4.0},
   year = {2020},
 }
 ```

flax/version.py

@@ -12,5 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version__ = "0.3.6"
+__version__ = "0.4.0"