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

google · Jan 27, 2022 · 695a880 · 695a880
2 parents 5de87fa + a6e7744
commit 695a880
Show file tree

Hide file tree

Showing 54 changed files with 309 additions and 7,438 deletions.
diff --git a/CHANGELOG.md b/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`.

diff --git a/docs/design_notes/arguments.md b/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.
diff --git a/docs/flax.nn.rst b/docs/flax.nn.rst