Version 0.6.4

What's Changed

• Improve and migrate model_surgery guide by @IvyZX in #2668
• Add deprecation warnings and docstring pointers to flax.training.lr_schedule by @IvyZX in #2702
• Introduce flax.io.NotFoundError to remove tensorflow dependency by @wookayin in #2697
• wrap Module properties by @cgarciae in #2541
• Update the checkpointing guide with Orbax migration code examples by @IvyZX in #2688
• Include mypy in run_all_tests by @zaxtax in #2670
• Add serialization for Partial functions. Completes #2433 by @zaxtax in #2557
• Cleanup DenseGeneral by @cgarciae in #2701
• Adds wave2vec2 link by @marcvanzee in #2718
• Remove some filterwarnings by @marcvanzee in #2715
• Adds standardize initializer by @marcvanzee in #2717
• Bump action/setup-python Python version to 3.9 by @marcvanzee in #2726
• Expose module.merge_param in RTD and fix docstring by @marcvanzee in #2680
• Revert to using public github runner pool while internal issues are fixed. by @levskaya in #2729
• Re-enable faster github CI runners by @levskaya in #2732
• Monitor checkpoint load/save durations. by @copybara-service in #2629
• Add Flax Dropout guide by @8bitmp3 in #2675
• Update readthedocs build to python 3.8 by @IvyZX in #2746
• Update Flax BatchNorm guide title by @8bitmp3 in #2747
• Simplify _AxisRules by @cgarciae in #2734
• Re-added serialization error checking for dicts by @chiamp in #2691
• Guide for Flax partitioning (auto SPMD) API by @IvyZX in #2730
• Fix LogicalNames type by @cgarciae in #2759
• removed duplicate flip section in docs by @chiamp in #2757
• updated tabulate documentation by @chiamp in #2756
• add unbind by @cgarciae in #2674
• Remove incorrect SPMD section in linen.rst by @cgarciae in #2766
• Adds some more comments to scope.py by @marcvanzee in #2743
• Fix GRU typo by @jheek in #2785
• Add metrics for async checkpoint (beacon and total write duration). by @copybara-service in #2786
• Fix indentation by @8bitmp3 in #2787
• Link CleanRL's PPO Implementation by @vwxyzjn in #2780
• Flax docs restructuring by @8bitmp3 in #2788
• Add metadata helper transform for adding axis metadata. by @levskaya in #2783
• Update Flax mission, fix some links by @8bitmp3 in #2797
• Updated getting started doc by @chiamp in #2698
• Added builder functions for zeros and ones initializers by @chiamp in #2790
• Restructure Flax Examples, add sections by @8bitmp3 in #2801
• Update how to build flax docs by @8bitmp3 in #2740
• Updating the doc by @arcAman07 in #2796
• added zeros_init and ones_init builder to initializers list by @chiamp in #2806
• Implement perturb no-op behaviour by @cgarciae in #2728
• Fix links in Flax Optax upgrading and dataset processing by @8bitmp3 in #2798
• improve Module.init docs by @cgarciae in #2792
• replaced zeros and ones initializers with zeros_init and ones_init builder functions by @chiamp in #2807
• Restructure Flax guides, add sections, update copyright year by @8bitmp3 in #2800
• Added glossary doc by @chiamp in #2805
• Add a comment about an error. by @nouiz in #2810
• Add transpose_kernel argument to ConvTranspose. by @mathisgerdes in #2578
• rename ConvLSTM to ConvLSTMCell by @cgarciae in #2767
• Add missing fields from repr by @cgarciae in #2803
• Allow specifying method as a string by @cgarciae in #2809
• Release 0.6.4 by @IvyZX in #2820

New Contributors

• @wookayin made their first contribution in #2697
• @vwxyzjn made their first contribution in #2780
• @arcAman07 made their first contribution in #2796
• @nouiz made their first contribution in #2810
• @mathisgerdes made their first contribution in #2578

Full Changelog: v0.6.3...v0.6.4