Merge pull request #1551 from jheek:version-0.3.5

PiperOrigin-RevId: 397933286

CHANGELOG.md

@@ -5,31 +5,58 @@ vNext
 ------
 
 (Add your change to a random empty line to avoid merge conflicts)
- -
- -
+- 
+- 
+- 
+- 
+- 
+- 
+- 
+- 
+- 
+- 
+- 
+- 
+- 
+- 
+- 
+- 
+- 
+- 
+- 
+- 
+- 
+- 
+- 
+- 
+
+
+
+0.3.5
+------
+
+Breaking changes:
+ - You can no longer pass an int as the `kernel_size` for a `flax.linen.Conv.
+   Instead a type error is raised stating that
+   a tuple/list should be provided. Stride and dilation arguments do support broadcasting a single int value now because this is not
+   ambigious when the kernel rank is known. 
+ - `flax.linen.enable_named_call` and `flax.linen.disable_named_call` now work anywhere instead of only affecting Modules constructed after the enable/disable call. Additionally, there is now `flax.linen.override_named_call` that provided a context manager to locally disable/enable named_call.
+ - NamedTuples are no longer converted to tuples on assignment to a `linen.Module`.
+
+New features:
+ - Flax internal stack frames are now removed from exception state traces.
+ - Added `flax.linen.nowrap` to decorate method that should not be transformed
+   because they are stateful.
  - Flax no longer uses implicit rank broadcasting. Thus, you can now use Flax with `--jax_numpy_rank_promotion=raise`.
- -
- -
- -
- - NamedTuples are no longer converted to tuples on assingment to a `linen.Module`.
- -
- -
+
+Bugfixes:
+ - linen Modules and dataclasses made with `flax.struct.dataclass` or `flax.struct.PyTreeNode` are now correctly recognized as dataclasses by static analysis tools like PyLance. Autocomplete of constructors has been verified to work with VSCode.
  - Fixed a bug in FrozenDict which didn't allow copying dicts with reserved names.
- -
- -
- -
- - `flax.linen.enable_named_call` and `flax.linen.disable_named_call` now work anywhere instead of only affecting Modules constructed after the enable/disable call. Additionally, there is now `flax.linen.override_named_call` that provided a context manager to locally disable/enable named_call.
- -
- -
  - Fix the serialization of named tuples. Tuple fields are no longer stored in the state dict and the named tuple class is no longer recreated ([bug](https://github.com/google/flax/issues/1429)).
- -
- - linen Modules and dataclasses made with `flax.struct.dataclass` or `flax.struct.PyTreeNode` are now correctly recognized as dataclasses by static analysis tools like PyLance. Autocomplete of constructors has been verified to work with VSCode.
- -
- -
- - `flax.linen.Conv` no longer interprets an int past as kernel_size as a 1d convolution. Instead a type error is raised stating that
-   a tuple/list should be provided. Stride and dilation arguments do support broadcasting a single int value now because this is not
-   ambigious when the kernel rank is known. 
- - 
+ - Mixed precision training with float16 now works correctly with the attention layers.
+ - auto-generated linen Module `__hash__`, `__eq__`, `__repr__` no longer fail by default on non-init attributes.
+
+
 
 0.3.4
 ------

README.md

@@ -179,7 +179,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.4},
+  version = {0.3.5},
   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.4"
+__version__ = "0.3.5"