Merge branch 'main' of https://www.github.com/google/flax into pipy

Author: marcvanzee

CHANGELOG.md

@@ -33,6 +33,8 @@ vNext
 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`.

docs/design_notes/arguments.md

@@ -85,9 +85,9 @@ It also avoids a default value which would probably cause either the train step
 
 
 
-## Functional Core and flax.nn
+## Functional Core
 
-The old NN api and functional core define functions rather than classes.
-Therefore, there is no clear distinction between hyper parameters and call time arguments.
-The only way to pre-determine the hyper parameters is by using `partial`.
+Functional core defines functions rather than classes.
+Therefore, there is no clear distinction between hyperparameters and call-time arguments.
+The only way to pre-determine the hyperparameters is by using `partial`.
 On the upside, there are no ambiguous cases where method arguments could also be attributes.