flax

What's Changed

• fixed codecov by @chiamp in https://github.com/google/flax/pull/3895
• Make FlatState a Mapping instead of a dict by @NeilGirdhar in https://github.com/google/flax/pull/3880
• Share nnx node registry between threads by @NeilGirdhar in https://github.com/google/flax/pull/3901
• fixed jnp.clip deprecation by @chiamp in https://github.com/google/flax/pull/3905
• Added three tab option to sphinx directive codediff and added testing for first tab by @chiamp in https://github.com/google/flax/pull/3847
• Add support for jax.sharding.PartitionSpec.UNCONSTRAINED in logical specification by @copybara-service in https://github.com/google/flax/pull/3902
• [nnx] fix mypy and pytype by @cgarciae in https://github.com/google/flax/pull/3894
• [nnx] fix iter_nodes by @cgarciae in https://github.com/google/flax/pull/3889
• [nnx] Sequential uses regular list by @cgarciae in https://github.com/google/flax/pull/3909
• [nnx] add ConvTranspose by @cgarciae in https://github.com/google/flax/pull/3908
• [nnx] add Module pytree_experimental static test by @cgarciae in https://github.com/google/flax/pull/3864
• Added docstring for Module.scope.path by @chiamp in https://github.com/google/flax/pull/3913
• [linen] test jit caching with state updates by @cgarciae in https://github.com/google/flax/pull/3900
• v0.8.4 by @cgarciae in https://github.com/google/flax/pull/3891
• [linen] enable separate initializers for out layer in MultiHeadDotProductAttention by @cgarciae in https://github.com/google/flax/pull/3835
• [nnx] cleanup graph by @cgarciae in https://github.com/google/flax/pull/3915
• [nnx] fix bugs by @cgarciae in https://github.com/google/flax/pull/3925
• Replace deprecated jax.tree_* functions with jax.tree.* by @copybara-service in https://github.com/google/flax/pull/3926
• [nnx] Object refactor by @cgarciae in https://github.com/google/flax/pull/3910
• [nnx] add iter_graph by @cgarciae in https://github.com/google/flax/pull/3919
• [nnx] add compat by @cgarciae in https://github.com/google/flax/pull/3921
• [nnx] transforms refactor by @cgarciae in https://github.com/google/flax/pull/3927
• added equivalence test for nnx.ConvTranspose by @chiamp in https://github.com/google/flax/pull/3934
• added equivalence test for nnx.Sequential by @chiamp in https://github.com/google/flax/pull/3935
• [NNX] Add LoRA and LoRALinear to NNX by @IvyZX in https://github.com/google/flax/pull/3929
• [nnx] fix substate mutability by @cgarciae in https://github.com/google/flax/pull/3932
• [nnx] improve update context by @cgarciae in https://github.com/google/flax/pull/3933
• [nnx] move out of experimental by @cgarciae in https://github.com/google/flax/pull/3936

Full Changelog: https://github.com/google/flax/compare/v0.8.3...v0.8.4

What's Changed

• Add git fetch upstream to contributing doc. by @carlosgmartin in https://github.com/google/flax/pull/3757
• removed getattr/setattr unboxing magic from nnx.Pytree by @chiamp in https://github.com/google/flax/pull/3743
• added Einsum layer to NNX by @chiamp in https://github.com/google/flax/pull/3741
• Make TrainState's step possibly jax.Array. This makes replicate valid for type checking. by @copybara-service in https://github.com/google/flax/pull/3763
• v0.8.3 by @cgarciae in https://github.com/google/flax/pull/3758
• [nnx] fix demo notebook by @cgarciae in https://github.com/google/flax/pull/3744
• added nnx api reference by @chiamp in https://github.com/google/flax/pull/3762
• updated rng docstring for init, apply and make_rng by @chiamp in https://github.com/google/flax/pull/3765
• use note box in make_rng docstring by @cgarciae in https://github.com/google/flax/pull/3767
• [nnx] improved graph update mechanism by @cgarciae in https://github.com/google/flax/pull/3759
• use note box in docstrings by @chiamp in https://github.com/google/flax/pull/3769
• Add reset_gate flag to MGUCell. by @carlosgmartin in https://github.com/google/flax/pull/3760
• Access thread_resources via jax.interpreters.pxla instead of jax.experimental.maps by @copybara-service in https://github.com/google/flax/pull/3775
• Minor doc improvements by @canyon289 in https://github.com/google/flax/pull/3588
• added MGU reset_gate test by @chiamp in https://github.com/google/flax/pull/3773
• [nnx] Pytrees are Trees by @cgarciae in https://github.com/google/flax/pull/3768
• Use short-circuiting access to debug_key_reuse by @copybara-service in https://github.com/google/flax/pull/3781
• fix tabulate on norm wrappers by @chiamp in https://github.com/google/flax/pull/3772
• Add kw_only struct.dataclass test by @chiamp in https://github.com/google/flax/pull/3651
• extended PyTreeNode to take dataclass kwargs by @chiamp in https://github.com/google/flax/pull/3785
• [nnx] Arrays are state by @cgarciae in https://github.com/google/flax/pull/3791
• [nnx] add GraphNode base class by @cgarciae in https://github.com/google/flax/pull/3790
• [nnx] jit accepts many Modules by @cgarciae in https://github.com/google/flax/pull/3783
• Exposing the experimental _split_transpose JAX scan parameter in Flax. by @copybara-service in https://github.com/google/flax/pull/3795
• Expose nnx.GraphNode by @chiamp in https://github.com/google/flax/pull/3796
• [nnx] Rngs and RngStream inherit from GraphNode by @cgarciae in https://github.com/google/flax/pull/3793
• [nnx] TrainState uses struct by @cgarciae in https://github.com/google/flax/pull/3788
• [nnx] split returns graphdef first by @cgarciae in https://github.com/google/flax/pull/3794
• Remove the uninitialized field "embedding" in nn.Embed by @copybara-service in https://github.com/google/flax/pull/3801
• Add nnx.training by @chiamp in https://github.com/google/flax/pull/3782
• [nnx] non-str State keys by @cgarciae in https://github.com/google/flax/pull/3802
• [nnx] allow all jit kwargs in nnx.jit by @cgarciae in https://github.com/google/flax/pull/3809
• [nnx] simplify readme by @cgarciae in https://github.com/google/flax/pull/3805
• [nnx] Fix nnx basics by @cgarciae in https://github.com/google/flax/pull/3812
• [nnx] grad accepts argnums by @cgarciae in https://github.com/google/flax/pull/3798
• [nnx] improve toy examples by @cgarciae in https://github.com/google/flax/pull/3813
• [nnx] expose Sequential by @cgarciae in https://github.com/google/flax/pull/3814
• [nnx] Rng Variable tags by @cgarciae in https://github.com/google/flax/pull/3807
• [nnx] remove copy in graph unflatten by @cgarciae in https://github.com/google/flax/pull/3804
• fixed optax guide links and docstring typos by @chiamp in https://github.com/google/flax/pull/3789
• added dropout broadcast test by @chiamp in https://github.com/google/flax/pull/3776
• relaxed grads kwarg for Optimizer.update by @chiamp in https://github.com/google/flax/pull/3818
• added tree_map deprecation warning filter by @chiamp in https://github.com/google/flax/pull/3828
• updated tree_map by @chiamp in https://github.com/google/flax/pull/3823
• added NNX vs JAX transformations guide by @chiamp in https://github.com/google/flax/pull/3819
• Updated NNX MNIST tutorial by @chiamp in https://github.com/google/flax/pull/3810
• [nnx] add Dropout.rngs by @cgarciae in https://github.com/google/flax/pull/3815
• removed autosummary from linen docs by @chiamp in https://github.com/google/flax/pull/3792
• Fix cloudpickle sentinel cloning by @cgarciae in https://github.com/google/flax/pull/3825
• [nnx] remove pytreelib by @cgarciae in https://github.com/google/flax/pull/3816
• [nnx] fix nnx_basics by @cgarciae in https://github.com/google/flax/pull/3839
• [linen] fix DenseGeneral init by @cgarciae in https://github.com/google/flax/pull/3834
• [nnx] jit constrain object state by @cgarciae in https://github.com/google/flax/pull/3817
• Copybara import of the project: by @copybara-service in https://github.com/google/flax/pull/3857
• Add example of unbox() and replace_boxed() to the jit guide by @IvyZX in https://github.com/google/flax/pull/3843
• RNNCellBase refactor FLIP by @cgarciae in https://github.com/google/flax/pull/3099
• [nnx] Some small documentation suggestions. by @gnecula in https://github.com/google/flax/pull/3861
• updated nnx dropout by @chiamp in https://github.com/google/flax/pull/3841
• Fix LogicalRules type annotation. (Tuple[str] is a tuple with single element string, by @copybara-service in https://github.com/google/flax/pull/3877
• Add option to skip float32 promotion when computing means and variances for normalization. by @copybara-service in https://github.com/google/flax/pull/3873
• added nnx api reference link by @chiamp in https://github.com/google/flax/pull/3871
• option of forcing the input of softmax to be fp32 for better numerical stability in mixed-precision training. by @copybara-service in https://github.com/google/flax/pull/3874
• allow custom dot_general for einsum. by @copybara-service in https://github.com/google/flax/pull/3884
• [NVIDIA] Extend the custom fp8 accumulate dtype in non-jit scenarios by @kaixih in https://github.com/google/flax/pull/3827
• updated robots.txt by @chiamp in https://github.com/google/flax/pull/3886
• fixed autosummary links by @chiamp in https://github.com/google/flax/pull/3887
• Fix jax.tree_util.register_dataclass in older JAX versions. by @copybara-service in https://github.com/google/flax/pull/3885
• [nnx] v0.1 by @cgarciae in https://github.com/google/flax/pull/3876

Full Changelog: https://github.com/google/flax/compare/v0.8.2...v0.8.3

What's Changed

• Add +1 to version after 0.8.1 release by @IvyZX in https://github.com/google/flax/pull/3684
• fixed rng guide outputs by @chiamp in https://github.com/google/flax/pull/3685
• enforce mask kwarg in norm layers by @chiamp in https://github.com/google/flax/pull/3663
• added kwargs to self.param and self.variable by @chiamp in https://github.com/google/flax/pull/3675
• added nnx normalization tests by @chiamp in https://github.com/google/flax/pull/3689
• added NNX init_cache docstring example by @chiamp in https://github.com/google/flax/pull/3688
• added nnx attention equivalence test by @chiamp in https://github.com/google/flax/pull/3687
• Fix bug that assumed frozen-dict keys were strings. by @copybara-service in https://github.com/google/flax/pull/3692
• added nnx rmsnorm by @chiamp in https://github.com/google/flax/pull/3691
• updated nnx compute_stats by @chiamp in https://github.com/google/flax/pull/3693
• fixed intercept_methods docstring by @chiamp in https://github.com/google/flax/pull/3694
• [nnx] Add Sphinx Docs by @cgarciae in https://github.com/google/flax/pull/3678
• Fix pointless docstring example of nn.checkpoint / nn.remat. by @levskaya in https://github.com/google/flax/pull/3703
• added default params rng to .apply by @chiamp in https://github.com/google/flax/pull/3698
• [nnx] add partial_init by @cgarciae in https://github.com/google/flax/pull/3674
• make make_rng default to 'params' by @chiamp in https://github.com/google/flax/pull/3699
• Add SimpleCell. by @carlosgmartin in https://github.com/google/flax/pull/3697
• fix Module.module_paths docstring by @cgarciae in https://github.com/google/flax/pull/3709
• Guarantee the latest JAX version on CI by @cgarciae in https://github.com/google/flax/pull/3705
• Replace deprecated API jax.tree_map by @copybara-service in https://github.com/google/flax/pull/3715
• Use jax.tree_util.tree_map instead of deprecated jax.tree_map. by @copybara-service in https://github.com/google/flax/pull/3714
• [nnx] simplify readme by @cgarciae in https://github.com/google/flax/pull/3707
• [nnx] add demo.ipynb by @cgarciae in https://github.com/google/flax/pull/3680
• Fix Tabulate's compute_flops by @cgarciae in https://github.com/google/flax/pull/3721
• [nnx] simplify TraceState by @cgarciae in https://github.com/google/flax/pull/3724
• Add broadcast of strides and kernel_dilation to nn.ConvTranspose by @IvyZX in https://github.com/google/flax/pull/3731
• [nnx] Fix State.sub by @cgarciae in https://github.com/google/flax/pull/3704
• [nnx] always fold_in on fork + new ForkedKeys return type by @cgarciae in https://github.com/google/flax/pull/3722
• [nnx] explicit Variables by @cgarciae in https://github.com/google/flax/pull/3720
• Improves fingerprint definition for Modules in nn.jit. by @copybara-service in https://github.com/google/flax/pull/3736
• Flax: avoid key reuse in tests by @copybara-service in https://github.com/google/flax/pull/3740
• added Einsum layer by @chiamp in https://github.com/google/flax/pull/3710
• nn.jit: automatic fingerprint definition for dataclass attributes by @cgarciae in https://github.com/google/flax/pull/3737
• [NVIDIA] Use custom grad accumulation for FP8 params by @kaixih in https://github.com/google/flax/pull/3623
• removed nnx dataclass by @chiamp in https://github.com/google/flax/pull/3742
• [nnx] cleanup graph_utils by @cgarciae in https://github.com/google/flax/pull/3728
• Fix doctest and unbreak head by @IvyZX in https://github.com/google/flax/pull/3753
• [nnx] add pytree support by @cgarciae in https://github.com/google/flax/pull/3732
• fixed intercept_methods docstring by @chiamp in https://github.com/google/flax/pull/3752
• Add ConvLSTMCell to docs. by @carlosgmartin in https://github.com/google/flax/pull/3712
• [nnx] remove flagslib by @cgarciae in https://github.com/google/flax/pull/3733
• Fix tests after applying JAX key-reuse checker. See: by @copybara-service in https://github.com/google/flax/pull/3748

Full Changelog: https://github.com/google/flax/compare/v0.8.1...v0.8.2

What's Changed

• bump version number to 0.8.1 by @chiamp in https://github.com/google/flax/pull/3649
• Bump pillow from 10.0.1 to 10.2.0 in /examples/vae by @dependabot in https://github.com/google/flax/pull/3641
• fixed docstring by @chiamp in https://github.com/google/flax/pull/3643
• Add explicit control over frozen/slots setting in flax.struct.dataclass by @copybara-service in https://github.com/google/flax/pull/3645
• make Sequential.call compact by @copybara-service in https://github.com/google/flax/pull/3647
• add Module.module_paths by @cgarciae in https://github.com/google/flax/pull/3654
• added rng_guide by @chiamp in https://github.com/google/flax/pull/3497
• Replacing jax.tree_util.tree_map with mapping over leafs. by @copybara-service in https://github.com/google/flax/pull/3658
• Copybara import of the project: by @copybara-service in https://github.com/google/flax/pull/3659
• added InstanceNorm by @chiamp in https://github.com/google/flax/pull/3652
• add Module.module_paths by @copybara-service in https://github.com/google/flax/pull/3660
• added norm equivalence tests by @chiamp in https://github.com/google/flax/pull/3662
• updated nowrap docstring by @chiamp in https://github.com/google/flax/pull/3661
• Add module_paths method to docs by @cgarciae in https://github.com/google/flax/pull/3657
• add default make_rng by @chiamp in https://github.com/google/flax/pull/3669
• renamed channel_axes to feature_axes in InstanceNorm by @chiamp in https://github.com/google/flax/pull/3667
• added flax.typing by @chiamp in https://github.com/google/flax/pull/3624
• changed kwargs to actual key-word args by @chiamp in https://github.com/google/flax/pull/3562
• updated docs and docstrings by @chiamp in https://github.com/google/flax/pull/3670
• re-added linen_intro by @chiamp in https://github.com/google/flax/pull/3672
• add compact_name_scope v3 by @cgarciae in https://github.com/google/flax/pull/3646
• Release 0.8.1 by @IvyZX in https://github.com/google/flax/pull/3682

Full Changelog: https://github.com/google/flax/compare/v0.8.0...v0.8.1

What's Changed

• bump version number by @levskaya in https://github.com/google/flax/pull/3446
• Add merge / finalize step when using OCDBT driver. Files will be first written to per-process subdirectories, which are later copied by reference to the main directory before the checkpoint is finalized. by @copybara-service in https://github.com/google/flax/pull/3426
• fixed quickstart by @chiamp in https://github.com/google/flax/pull/3451
• [NVIDIA] Update the algorithm to compute fp8 scales by @kaixih in https://github.com/google/flax/pull/3441
• added pre-commit hook that sort imports and formats by @chiamp in https://github.com/google/flax/pull/3455
• restructured doc folders by @chiamp in https://github.com/google/flax/pull/3434
• Forked a subset of JAX configuration APIs by @superbobry in https://github.com/google/flax/pull/3448
• Fix Module.clone in deepclone mode for internal usage. by @levskaya in https://github.com/google/flax/pull/3459
• Add user-friendly module copy method. by @levskaya in https://github.com/google/flax/pull/3461
• Add simple argument-only lifted nn.grad function. by @levskaya in https://github.com/google/flax/pull/3463
• exempt a jax.config deprecation warning by @levskaya in https://github.com/google/flax/pull/3465
• Clean up pyproject.toml. by @levskaya in https://github.com/google/flax/pull/3468
• Allow for fast accumulation selection for FP8 GEMM by @wenscarl in https://github.com/google/flax/pull/3416
• re-added quickstart guide by @chiamp in https://github.com/google/flax/pull/3471
• fixed tabulate docstring by @chiamp in https://github.com/google/flax/pull/3452
• Add NNX by @cgarciae in https://github.com/google/flax/pull/3218
• Bump pillow from 9.5.0 to 10.0.1 in /examples/vae by @dependabot in https://github.com/google/flax/pull/3390
• updated attention_test by @chiamp in https://github.com/google/flax/pull/3454
• [nnx] Improve docs by @cgarciae in https://github.com/google/flax/pull/3478
• added example docstrings by @chiamp in https://github.com/google/flax/pull/3453
• fix nn.value_and_grad by implementing directly in core by @levskaya in https://github.com/google/flax/pull/3479
• Add dataset loading guide (Issue #2116) by @VictorPrins in https://github.com/google/flax/pull/3450
• [nnx] Add support for python container types by @cgarciae in https://github.com/google/flax/pull/3486
• remove SelfAttention test and warning filter by @chiamp in https://github.com/google/flax/pull/3470
• disabled ruff formatter by @chiamp in https://github.com/google/flax/pull/3482
• adding doctest to .rst files by @chiamp in https://github.com/google/flax/pull/3481
• changed pip installs to use quotes by @chiamp in https://github.com/google/flax/pull/3477
• added enum support for tabulate by @chiamp in https://github.com/google/flax/pull/3485
• fix bug in optimizer-api.md by @zhaoyang-0204 in https://github.com/google/flax/pull/3462
• removed selfattention from doctest by @chiamp in https://github.com/google/flax/pull/3489
• [nnx] Add missing import on why.ipynb by @cgarciae in https://github.com/google/flax/pull/3503
• [nnx] switch to nested State representation by @cgarciae in https://github.com/google/flax/pull/3502
• Improved Rigor of PReLU Test by @Micky774 in https://github.com/google/flax/pull/3498
• added geglu activation and tests by @HMUNACHI in https://github.com/google/flax/pull/3512
• [nnx] Add LinearGeneral and MultiHeadAttention by @cgarciae in https://github.com/google/flax/pull/3487
• Add NNX/Linen consistency test for Embed layer by @Micky774 in https://github.com/google/flax/pull/3513
• Add NNX/Linen API consistency test for Conv layer by @Micky774 in https://github.com/google/flax/pull/3511
• Prevent crash in dataclasses with no-init params by @NeilGirdhar in https://github.com/google/flax/pull/3514
• [nnx] Variable referece sharing by @cgarciae in https://github.com/google/flax/pull/3516
• Added NNX/Linen API consistency test for Linear/Dense layer by @Micky774 in https://github.com/google/flax/pull/3509
• Add missing mask argument to LayerNorm, RMSNorm, and GroupNorm. by @carlosgmartin in https://github.com/google/flax/pull/3510
• [nnx] Fix graph_utils bug by @cgarciae in https://github.com/google/flax/pull/3518
• remove deprecated normalize function by @chiamp in https://github.com/google/flax/pull/3531
• Reduced number of parameterizations for Conv NNX/Linen consistency test by @Micky774 in https://github.com/google/flax/pull/3526
• Ensure that _hashable_filter does not convert strings to a tuple of letters by @copybara-service in https://github.com/google/flax/pull/3533
• added sow attention weights by @chiamp in https://github.com/google/flax/pull/3529
• Fix scan out_axes by @cgarciae in https://github.com/google/flax/pull/3540
• updated embed docstring by @chiamp in https://github.com/google/flax/pull/3539
• add test_scan_negative_axes by @cgarciae in https://github.com/google/flax/pull/3542
• add module methods to api docs by @chiamp in https://github.com/google/flax/pull/3544
• fixed double backquote code font by @chiamp in https://github.com/google/flax/pull/3545
• add nnx conv support for int kernel size by @chiamp in https://github.com/google/flax/pull/3537
• added sow attention weights to NNX by @chiamp in https://github.com/google/flax/pull/3548
• changed return_weights to sow_weights for attention layer by @chiamp in https://github.com/google/flax/pull/3550
• format linen_linear_test.py by @chiamp in https://github.com/google/flax/pull/3553
• re-factored features arg by @chiamp in https://github.com/google/flax/pull/3554
• updated NNX readme by @chiamp in https://github.com/google/flax/pull/3556
• Disable ruff sort imports by @cgarciae in https://github.com/google/flax/pull/3560
• Add StateVariablesMapping by @cgarciae in https://github.com/google/flax/pull/3523
• add kwargs support for nn.jit by @copybara-service in https://github.com/google/flax/pull/3559
• [nnx] Fix readme install instruction by @cgarciae in https://github.com/google/flax/pull/3565
• implement Rng.getattr by @cgarciae in https://github.com/google/flax/pull/3547
• [nnx] add qkv_features back to MHA by @cgarciae in https://github.com/google/flax/pull/3566
• updated readme by @chiamp in https://github.com/google/flax/pull/3563
• fixed typo by @chiamp in https://github.com/google/flax/pull/3561
• Raise an error for a bad key type by @NeilGirdhar in https://github.com/google/flax/pull/3527
• re-factored nnx initializers by @chiamp in https://github.com/google/flax/pull/3555
• [nnx] Add complex test with scan + batchnorm + dropout by @cgarciae in https://github.com/google/flax/pull/3567
• [nnx] Add interacting with JAX section to README by @cgarciae in https://github.com/google/flax/pull/3573
• expose ones and zeros initializers by @chiamp in https://github.com/google/flax/pull/3574
• Fix promotion bug in MultiHeadDotProductAttention: by @giovannic in https://github.com/google/flax/pull/3571
• fixed error doc formatting by @chiamp in https://github.com/google/flax/pull/3587
• [nnx] Improve spmd by @cgarciae in https://github.com/google/flax/pull/3580
• [nnx] improve graph_utils._set_key_tuple by @cgarciae in https://github.com/google/flax/pull/3592
• [nnx] Fix variable unflatten by @cgarciae in https://github.com/google/flax/pull/3578
• [nnx] add open in colab button to why nnx by @cgarciae in https://github.com/google/flax/pull/3596
• [nnx] Export missing symbols by @cgarciae in https://github.com/google/flax/pull/3583
• [nnx] flaglib add get overloads by @cgarciae in https://github.com/google/flax/pull/3582
• Fix type in NNX readme by @shoyer in https://github.com/google/flax/pull/3591
• [nnx] add submodule iterator by @cgarciae in https://github.com/google/flax/pull/3581
• [nnx] delete flaglib duplicated copyright comment by @cgarciae in https://github.com/google/flax/pull/3600
• fixed NNX decode and dynamic slicing by @chiamp in https://github.com/google/flax/pull/3576
• [nnx] cleanup CallableProxy by @cgarciae in https://github.com/google/flax/pull/3608
• [nnx] improve runtime flags by @cgarciae in https://github.com/google/flax/pull/3607
• fixed broken links on quick start guide by @chiamp in https://github.com/google/flax/pull/3610
• added multiheadattention alias by @chiamp in https://github.com/google/flax/pull/3572
• Rollback of Copybara import of the project: by @copybara-service in https://github.com/google/flax/pull/3612
• add missing docs for module functions by @cgarciae in https://github.com/google/flax/pull/3619
• fix lm1b data sharding by @cgarciae in https://github.com/google/flax/pull/3620
• improve embed by @jianyizh in https://github.com/google/flax/pull/3590
• disable ruff linter by @chiamp in https://github.com/google/flax/pull/3625
• Add compact_name_scope decorator by @cgarciae in https://github.com/google/flax/pull/3621
• Copybara import of the project: by @copybara-service in https://github.com/google/flax/pull/3638
• added BatchApply by @chiamp in https://github.com/google/flax/pull/3634
• add compact_name_scope v2 by @copybara-service in https://github.com/google/flax/pull/3640
• add compact_name_scope v2 by @copybara-service in https://github.com/google/flax/pull/3642
• release 0.8.0 by @chiamp in https://github.com/google/flax/pull/3644

New Contributors

• @superbobry made their first contribution in https://github.com/google/flax/pull/3448
• @VictorPrins made their first contribution in https://github.com/google/flax/pull/3450
• @zhaoyang-0204 made their first contribution in https://github.com/google/flax/pull/3462
• @Micky774 made their first contribution in https://github.com/google/flax/pull/3498
• @HMUNACHI made their first contribution in https://github.com/google/flax/pull/3512
• @carlosgmartin made their first contribution in https://github.com/google/flax/pull/3510
• @giovannic made their first contribution in https://github.com/google/flax/pull/3571
• @jianyizh made their first contribution in https://github.com/google/flax/pull/3590

Full Changelog: https://github.com/google/flax/compare/v0.7.5...v0.8.0

What's Changed

• Add method-to-model section to Haiku migration guide by @IvyZX in https://github.com/google/flax/pull/3277
• updated haiku guide with new JAX RNG api by @chiamp in https://github.com/google/flax/pull/3343
• changed resnet v1 to v1.5 by @chiamp in https://github.com/google/flax/pull/3344
• updated flax_basics by @chiamp in https://github.com/google/flax/pull/3342
• Add find methods and magic methods for Cursor API by @chiamp in https://github.com/google/flax/pull/3306
• fix DeprecationWarnings by @chiamp in https://github.com/google/flax/pull/3352
• Make checkpoint guide path absolute by @IvyZX in https://github.com/google/flax/pull/3358
• use jax.Array type for rng keys by @chiamp in https://github.com/google/flax/pull/3354
• Disable pyink by @cgarciae in https://github.com/google/flax/pull/3356
• removed DeprecationWarning filter by @chiamp in https://github.com/google/flax/pull/3359
• Don't propagate default args in Tabulate by @cgarciae in https://github.com/google/flax/pull/3357
• added spectral norm by @chiamp in https://github.com/google/flax/pull/3335
• fixed bind-unbind bug by @chiamp in https://github.com/google/flax/pull/3365
• Add MaxText open source LLM to index.rst by @8bitmp3 in https://github.com/google/flax/pull/3368
• fixed typo by @chiamp in https://github.com/google/flax/pull/3367
• Add fp8 custom op and unit test by @wenscarl in https://github.com/google/flax/pull/3322
• fixed scope typing by @chiamp in https://github.com/google/flax/pull/3371
• Trailing whitespace fixes. by @levskaya in https://github.com/google/flax/pull/3373
• Make fp8 ops use explicit broadcasting. by @levskaya in https://github.com/google/flax/pull/3374
• Conditonal contraints for numpy and clu by @cgarciae in https://github.com/google/flax/pull/3394
• add truncated_normal to initializers by @levskaya in https://github.com/google/flax/pull/3401
• fix HEAD by @chiamp in https://github.com/google/flax/pull/3404
• updated imagenet readme by @chiamp in https://github.com/google/flax/pull/3383
• Add Flax FAQ - how to search, how to take the derivative w.r.t. a hidden layer, remat_scan vs scan(remat), recommended training libraries/metrics by @8bitmp3 in https://github.com/google/flax/pull/3267
• added rmsnorm to api docs by @chiamp in https://github.com/google/flax/pull/3406
• updated docs by @chiamp in https://github.com/google/flax/pull/3370
• fix HEAD by @chiamp in https://github.com/google/flax/pull/3408
• added weightnorm layer by @chiamp in https://github.com/google/flax/pull/3405
• added attention refactor to changelog by @chiamp in https://github.com/google/flax/pull/3412
• added dropout arg to MultiHeadDotProductAttention by @chiamp in https://github.com/google/flax/pull/3384
• fix spectralnorm layer by @chiamp in https://github.com/google/flax/pull/3403
• remove pdf target by @cgarciae in https://github.com/google/flax/pull/3415
• added import precommit hook by @chiamp in https://github.com/google/flax/pull/3410
• fixed GRU docstring by @chiamp in https://github.com/google/flax/pull/3419
• Replaces pjit with jit in spmd.py by @copybara-service in https://github.com/google/flax/pull/3421
• Ignore transient chex deprecationwarning. by @levskaya in https://github.com/google/flax/pull/3427
• Remove transformers dependency from docs by @cgarciae in https://github.com/google/flax/pull/3431
• updated pytorch upgrade guide by @chiamp in https://github.com/google/flax/pull/3432
• Simplify abstract rng creation in param shape check. by @levskaya in https://github.com/google/flax/pull/3429
• [NVIDIA] Update the FP8 support by @kaixih in https://github.com/google/flax/pull/3435
• fixed mypy errors at HEAD by @chiamp in https://github.com/google/flax/pull/3440
• added has_improved field to EarlyStopping by @chiamp in https://github.com/google/flax/pull/3385
• fully deprecated old RNN api by @chiamp in https://github.com/google/flax/pull/3425
• updated lm1b example with jit by @chiamp in https://github.com/google/flax/pull/3302
• added MGU class by @chiamp in https://github.com/google/flax/pull/3418
• Make Flax Basics visible by @8bitmp3 in https://github.com/google/flax/pull/3443
• update for release v0.7.5 by @levskaya in https://github.com/google/flax/pull/3444

New Contributors

• @wenscarl made their first contribution in https://github.com/google/flax/pull/3322
• @kaixih made their first contribution in https://github.com/google/flax/pull/3435

Full Changelog: https://github.com/google/flax/compare/v0.7.4...v0.7.5

What's Changed

Added python version constraint >=3.9.

• Bump version no by 1 post-release. by @copybara-service in https://github.com/google/flax/pull/3328
• Update minimum python version info. by @levskaya in https://github.com/google/flax/pull/3331

Full Changelog: https://github.com/google/flax/compare/v0.7.3...v0.7.4

What's Changed

• Fix Documentation Typo by @peterdavidfagan in https://github.com/google/flax/pull/3265
• bump version number by @chiamp in https://github.com/google/flax/pull/3273
• added cursor api by @chiamp in https://github.com/google/flax/pull/3246
• Improve Typing support by @cgarciae in https://github.com/google/flax/pull/3242
• feat: add configs for vae example by @SauravMaheshkar in https://github.com/google/flax/pull/3254
• fix stackoverflow when loading pickled module by @cgarciae in https://github.com/google/flax/pull/3286
• Remove string from checkpointing example and unbreak doctest in head by @IvyZX in https://github.com/google/flax/pull/3288
• Improve kw_only_dataclass by @cgarciae in https://github.com/google/flax/pull/3293
• updated haiku upgrade guide by @chiamp in https://github.com/google/flax/pull/3292
• added logical partitioning to pjit guide by @chiamp in https://github.com/google/flax/pull/3290
• Rolling back google/flax@ba9e24a7b697e6407465cb4b05a24a4cea152248 by @copybara-service in https://github.com/google/flax/pull/3295
• Add RNN FLIP by @cgarciae in https://github.com/google/flax/pull/2585
• [flax] Add module path to nn.module. by @copybara-service in https://github.com/google/flax/pull/3300
• Minor fix to jax_utils.prefect_to_device by @anuragarnab in https://github.com/google/flax/pull/3308
• add tf-text to doc requirements by @cgarciae in https://github.com/google/flax/pull/3317
• Allow apply's method argument to accept submodules by @cgarciae in https://github.com/google/flax/pull/3281
• Update handling of typed PRNG keys by @jakevdp in https://github.com/google/flax/pull/3314
• 0.7.3 release by @IvyZX in https://github.com/google/flax/pull/3326

New Contributors

• @peterdavidfagan made their first contribution in https://github.com/google/flax/pull/3265
• @SauravMaheshkar made their first contribution in https://github.com/google/flax/pull/3254

Full Changelog: https://github.com/google/flax/compare/v0.7.2...v0.7.3

What's Changed

• make flax.core.copy add_or_replace optional by @PhilipVinc in https://github.com/google/flax/pull/3241
• bump version to 0.7.2 by @chiamp in https://github.com/google/flax/pull/3258
• Release 0.7.2 by @chiamp in https://github.com/google/flax/pull/3269

Full Changelog: https://github.com/google/flax/compare/v0.7.1...v0.7.2

What's Changed

• added dict migration guide to index by @chiamp in https://github.com/google/flax/pull/3188
• [linen] Add alternative, more numerically stable, variance calculation to LayerNorm. by @copybara-service in https://github.com/google/flax/pull/3194
• [linen] Minor cleanup to normalization code. by @copybara-service in https://github.com/google/flax/pull/3200
• Fix warnings from atari gym. by @levskaya in https://github.com/google/flax/pull/3207
• Fix carry slice logic by @cgarciae in https://github.com/google/flax/pull/3213
• removed FrozenDict section by @chiamp in https://github.com/google/flax/pull/3210
• make flax_basics guide use utility fns by @chiamp in https://github.com/google/flax/pull/3214
• Fix typo in struct.py documentation by @marcoselvi in https://github.com/google/flax/pull/3209
• updated flax guides to use utility fns by @chiamp in https://github.com/google/flax/pull/3215
• [linen] More minor cleanup in normalization compute_stats. by @copybara-service in https://github.com/google/flax/pull/3205
• Expose options to customize rich.Table by @cgarciae in https://github.com/google/flax/pull/3197
• Use pyink by @cgarciae in https://github.com/google/flax/pull/3216
• Fix remaining pyink issues by @cgarciae in https://github.com/google/flax/pull/3219
• add scan over layers section by @cgarciae in https://github.com/google/flax/pull/3195
• Fix checkpointing guide error at head by @IvyZX in https://github.com/google/flax/pull/3223
• Improve scan docs by @cgarciae in https://github.com/google/flax/pull/3231
• 0.7.1 release by @chiamp in https://github.com/google/flax/pull/3234

New Contributors

• @marcoselvi made their first contribution in https://github.com/google/flax/pull/3209

Full Changelog: https://github.com/google/flax/compare/v0.7.0...v0.7.1

What's Changed

• added dict migration guide to index by @chiamp in https://github.com/google/flax/pull/3188
• [linen] Add alternative, more numerically stable, variance calculation to LayerNorm. by @copybara-service in https://github.com/google/flax/pull/3194
• [linen] Minor cleanup to normalization code. by @copybara-service in https://github.com/google/flax/pull/3200
• Fix warnings from atari gym. by @levskaya in https://github.com/google/flax/pull/3207
• Fix carry slice logic by @cgarciae in https://github.com/google/flax/pull/3213
• removed FrozenDict section by @chiamp in https://github.com/google/flax/pull/3210
• make flax_basics guide use utility fns by @chiamp in https://github.com/google/flax/pull/3214
• Fix typo in struct.py documentation by @marcoselvi in https://github.com/google/flax/pull/3209
• updated flax guides to use utility fns by @chiamp in https://github.com/google/flax/pull/3215
• [linen] More minor cleanup in normalization compute_stats. by @copybara-service in https://github.com/google/flax/pull/3205
• Expose options to customize rich.Table by @cgarciae in https://github.com/google/flax/pull/3197
• Use pyink by @cgarciae in https://github.com/google/flax/pull/3216
• Fix remaining pyink issues by @cgarciae in https://github.com/google/flax/pull/3219
• add scan over layers section by @cgarciae in https://github.com/google/flax/pull/3195
• Fix checkpointing guide error at head by @IvyZX in https://github.com/google/flax/pull/3223
• Improve scan docs by @cgarciae in https://github.com/google/flax/pull/3231

New Contributors

• @marcoselvi made their first contribution in https://github.com/google/flax/pull/3209

Full Changelog: https://github.com/google/flax/compare/v0.7.0...v0.7.1

What's Changed

• delete long-unused dotgetter utility by @copybara-service in https://github.com/google/flax/pull/3156
• Add sections to flax.linen toctree by @cgarciae in https://github.com/google/flax/pull/3073
• update doctest requirements to use recentish jax minver by @levskaya in https://github.com/google/flax/pull/3160
• Update gym to gymnasium by @cgarciae in https://github.com/google/flax/pull/3133
• Add lifted transforms section to Haiku migration guide by @cgarciae in https://github.com/google/flax/pull/3158
• relaxed tolerance by @chiamp in https://github.com/google/flax/pull/3162
• fixed broken links by @chiamp in https://github.com/google/flax/pull/3161
• added dict migration guide by @chiamp in https://github.com/google/flax/pull/3109
• Update requirements, restructure files and fix formatting for VAE example by @canyon289 in https://github.com/google/flax/pull/3046
• Update python version support by @cgarciae in https://github.com/google/flax/pull/3168
• Improve Haiku migration guide by @cgarciae in https://github.com/google/flax/pull/3169
• Set default types in Flax for Orbax restoration and add restore_with_serialized_types in preparation for an upcoming change. by @copybara-service in https://github.com/google/flax/pull/3165
• fix precommit issues by @chiamp in https://github.com/google/flax/pull/3170
• fixed incorrect reference link by @chiamp in https://github.com/google/flax/pull/3167
• added absltest to linen_recurrent_test by @chiamp in https://github.com/google/flax/pull/3172
• remove cell_size docs from RNN by @cgarciae in https://github.com/google/flax/pull/3186
• 0.7.0 by @cgarciae in https://github.com/google/flax/pull/3187

Full Changelog: https://github.com/google/flax/compare/v0.6.11...v0.7.0

What's Changed

• remove unreferenced docs by @chiamp in https://github.com/google/flax/pull/3085
• Bump tensorflow from 2.9.3 to 2.11.1 in /examples/ogbg_molpcba by @dependabot in https://github.com/google/flax/pull/2976
• Bind submodules using getattribute by @cgarciae in https://github.com/google/flax/pull/2901
• Allow manually specifying the rng key for Dropout by @cgarciae in https://github.com/google/flax/pull/3114
• fixed missing separator for rng fold in by @chiamp in https://github.com/google/flax/pull/3095
• Bump imagenet Flax to 0.6.5 by @IvyZX in https://github.com/google/flax/pull/3118
• Rolling back #3104 because internal tests are breaking. by @copybara-service in https://github.com/google/flax/pull/3126
• Use pyproject.toml as hash on CI by @cgarciae in https://github.com/google/flax/pull/3131
• Minor improvements for Sequential by @cgarciae in https://github.com/google/flax/pull/3129
• Update activation.py docstring for PReLU by @levskaya in https://github.com/google/flax/pull/3122
• Eagerly bind submodules by @cgarciae in https://github.com/google/flax/pull/3077
• Add haiku_migration_guide by @cgarciae in https://github.com/google/flax/pull/3139
• Update Haiku migration guide. by @LenaMartens in https://github.com/google/flax/pull/3154
• Fix setup trigger issue under sharing and transforms. by @levskaya in https://github.com/google/flax/pull/3153
• Upgrade the pjit guide to jax.jit by @IvyZX in https://github.com/google/flax/pull/3016
• Release 0.6.11 by @IvyZX in https://github.com/google/flax/pull/3155

New Contributors

• @LenaMartens made their first contribution in https://github.com/google/flax/pull/3154

Full Changelog: https://github.com/google/flax/compare/v0.6.10...v0.6.11

What's Changed

• Add missing Bidirectional api reference by @cgarciae in https://github.com/google/flax/pull/2999
• Optimize struct.dataclass by @cgarciae in https://github.com/google/flax/pull/3024
• Rectified typos & grammar in docs by @Neilblaze in https://github.com/google/flax/pull/3021
• Updated full_eval.ipynb & corrected to constructor by @Neilblaze in https://github.com/google/flax/pull/3059
• fixed broken coverage link by @chiamp in https://github.com/google/flax/pull/3058
• Allow clipping the dynamic loss scale to a minimum value by @copybara-service in https://github.com/google/flax/pull/3051
• updated extracting_intermediates doc by @chiamp in https://github.com/google/flax/pull/3063
• updated api reference by @chiamp in https://github.com/google/flax/pull/3068
• fix RTD urllib3 error by @chiamp in https://github.com/google/flax/pull/3081
• Add test case for degenerate and deterministic attention layer dropout. by @copybara-service in https://github.com/google/flax/pull/3083
• 0.6.10 by @cgarciae in https://github.com/google/flax/pull/3089

New Contributors

• @Neilblaze made their first contribution in https://github.com/google/flax/pull/3021

Full Changelog: https://github.com/google/flax/compare/v0.6.9...v0.6.10

What's Changed

• Follow-up to #2985 (adding config.prefetch) by @copybara-service in https://github.com/google/flax/pull/3004
• Bringing changes in https://github.com/google/flax/pull/2955 internally. by @copybara-service in https://github.com/google/flax/pull/2960
• Updates imagenet.ipynb to use GPU Colab runtime. by @andsteing in https://github.com/google/flax/pull/3007
• Simplify pyproject.toml's mypy section by @cgarciae in https://github.com/google/flax/pull/3011
• Invalidate CI cache daily by @cgarciae in https://github.com/google/flax/pull/3010
• Update module.py by @aniquetahir in https://github.com/google/flax/pull/3006
• Fix get_partition_spec on replicated array by @IvyZX in https://github.com/google/flax/pull/3013
• Move setup.py to project.toml by @cgarciae in https://github.com/google/flax/pull/3012
• modify git repo master to main by @Mddct in https://github.com/google/flax/pull/3020
• Added pretty_repr utility fn by @chiamp in https://github.com/google/flax/pull/3022
• Update checkpointing guide with Orbax APIs by @IvyZX in https://github.com/google/flax/pull/3002
• Checkpoint fixes to unbreak head by @IvyZX in https://github.com/google/flax/pull/3027
• Use backend dependent os.path.getsize. by @copybara-service in https://github.com/google/flax/pull/3028
• Add more explicit warnings and logs for code with Orbax backend. by @copybara-service in https://github.com/google/flax/pull/3030
• Update orbax_upgrade_guide.rst fixing broken link by @canyon289 in https://github.com/google/flax/pull/3031
• Back to orbax-checkpoint after the package is fixed in 0.2.1. by @IvyZX in https://github.com/google/flax/pull/3029
• 0.6.9 release by @IvyZX in https://github.com/google/flax/pull/3035
• Fix release CI script by @cgarciae in https://github.com/google/flax/pull/3039

New Contributors

• @aniquetahir made their first contribution in https://github.com/google/flax/pull/3006
• @Mddct made their first contribution in https://github.com/google/flax/pull/3020
• @canyon289 made their first contribution in https://github.com/google/flax/pull/3031

Full Changelog: https://github.com/google/flax/compare/v0.6.8...v0.6.9

Changelog

• The automatic checkpoint migration was temporarily rolled back due to legacy compatibility issues.
  • We still recommend you to use the upgrade guide and migrate completely to the Orbax API to ensure stability.
  • Or alternatively, add flax.config.update('flax_use_orbax_checkpointing', True) to your project to avoid being impacted by the automatic migration process.
• Added utility functions to frozen_dict api.
• Migrated Flax away from register_keypaths.
• Fixes kwargs in convert_to_graphs_tuple_fn.
• Fixed examples in a few ways:
  • Bumped the TF version
  • Used latest checkpoint formats
  • Other misc fixes.

List of PRs

• Add Google Research examples with Flax by @8bitmp3 in https://github.com/google/flax/pull/2934
• [Trivial] [Post-release] Bump up the HEAD version number by @IvyZX in https://github.com/google/flax/pull/2946
• fix tabulate duplicating parameters by @cgarciae in https://github.com/google/flax/pull/2931
• permit passing mesh or rules to spmd primitives by @cgarciae in https://github.com/google/flax/pull/2939
• fixed ScopeParamShapeError message by @chiamp in https://github.com/google/flax/pull/2947
• Fix checkpointing docs for Orbax 0.1.4 by @IvyZX in https://github.com/google/flax/pull/2962
• Migrate Flax away from register_keypaths by @copybara-service in https://github.com/google/flax/pull/2963
• set axis='batch' by @yechengxi in https://github.com/google/flax/pull/2967
• Add Bidirectional by @cgarciae in https://github.com/google/flax/pull/2770
• Bump tensorflow from 2.9.3 to 2.11.1 in /examples/ppo by @dependabot in https://github.com/google/flax/pull/2984
• Bump tensorflow from 2.9.3 to 2.11.1 in /examples/imagenet by @dependabot in https://github.com/google/flax/pull/2983
• Bump tensorflow from 2.9.3 to 2.11.1 in /examples/sst2 by @dependabot in https://github.com/google/flax/pull/2980
• Bump tensorflow from 2.9.3 to 2.11.1 in /examples/nlp_seq by @dependabot in https://github.com/google/flax/pull/2979
• Bump tensorflow from 2.9.3 to 2.11.1 in /examples/lm1b by @dependabot in https://github.com/google/flax/pull/2977
• Adds missing api references by @cgarciae in https://github.com/google/flax/pull/2926
• Added utility functions to frozen_dict api by @chiamp in https://github.com/google/flax/pull/2975
• Fixed pop utility fn doc string by @chiamp in https://github.com/google/flax/pull/2958
• release 0.6.8 by @IvyZX in https://github.com/google/flax/pull/2987

Full Changelog: https://github.com/google/flax/compare/v0.6.7...v0.6.8

What's Changed

• Fix lm1b sampler by @cgarciae in https://github.com/google/flax/pull/2842
• RNN Module by @cgarciae in https://github.com/google/flax/pull/2604
• Add Orbax migration guide by @IvyZX in https://github.com/google/flax/pull/2885
• Add 'all' install group by @cgarciae in https://github.com/google/flax/pull/2894
• Improve nn.scan docs by @cgarciae in https://github.com/google/flax/pull/2839
• make utility functions work on FrozenDicts and regular dicts by @chiamp in https://github.com/google/flax/pull/2919
• Improve Flax dark mode - remove hard-coded text color by @8bitmp3 in https://github.com/google/flax/pull/2924
• Fix pytype and improve typing for Module by @cgarciae in https://github.com/google/flax/pull/2933
• Make flax.version 1 + the current release by @IvyZX in https://github.com/google/flax/pull/2936
• Migrate Flax dataclass to the newest JAX pytree keypath API. by @copybara-service in https://github.com/google/flax/pull/2921
• release 0.6.7 by @IvyZX in https://github.com/google/flax/pull/2944

Full Changelog: https://github.com/google/flax/compare/v0.6.6...v0.6.7

What's Changed

• Enhance checkpointing docs with Orbax information by @cpgaffney1 in https://github.com/google/flax/pull/2879
• A small doc and one fixe to Imagenet benchmark by @nouiz in https://github.com/google/flax/pull/2812
• Refactor and separate out name relaxation policy changes. by @levskaya in https://github.com/google/flax/pull/2873
• Added RMS normalization layer by @chiamp in https://github.com/google/flax/pull/2881
• added feature flag for FrozenDicts by @chiamp in https://github.com/google/flax/pull/2872
• Explicitly mark the lazy-rng with legacy_config, and all new configs should use flax.config. by @copybara-service in https://github.com/google/flax/pull/2895
• Added test for FrozenDict feature flag by @chiamp in https://github.com/google/flax/pull/2896
• Added dict utility functions by @chiamp in https://github.com/google/flax/pull/2897
• Improve notebook handling by @cgarciae in https://github.com/google/flax/pull/2860
• migrate tests to regular dicts by @chiamp in https://github.com/google/flax/pull/2898
• Release 0.6.6 by @IvyZX in https://github.com/google/flax/pull/2908

New Contributors

• @cpgaffney1 made their first contribution in https://github.com/google/flax/pull/2879

Full Changelog: https://github.com/google/flax/compare/v0.6.5...v0.6.6

What's Changed

• Internal change by @copybara-service in https://github.com/google/flax/pull/2826
• Cache .mypy_cache on CI by @cgarciae in https://github.com/google/flax/pull/2819
• improve nn.map_variables docstrings by @cgarciae in https://github.com/google/flax/pull/2818
• Render non-array variables in tabulate by @cgarciae in https://github.com/google/flax/pull/2769
• Improve nn.map_variables example by @cgarciae in https://github.com/google/flax/pull/2831
• Introduce relaxed naming policy behind a feature flag. by @levskaya in https://github.com/google/flax/pull/2833
• Bumps minimal JAX version from 0.3.16 to 0.3.24 and cleans up some code. by @copybara-service in https://github.com/google/flax/pull/2840
• import initializers module directly by @chiamp in https://github.com/google/flax/pull/2830
• Re-box axis metadata in DenseGeneral's init functions by @IvyZX in https://github.com/google/flax/pull/2843
• Update jax dependency by @jheek in https://github.com/google/flax/pull/2827
• replaced zeros and ones initializers with zeros_init and ones_init builder functions by @chiamp in https://github.com/google/flax/pull/2815
• Fixing the punctuation in the README by @arcAman07 in https://github.com/google/flax/pull/2811
• Implement lazy init by @jheek in https://github.com/google/flax/pull/2816
• improve nn.checkpoint docs by @cgarciae in https://github.com/google/flax/pull/2837
• Flax 0.6.5 point release by @levskaya in https://github.com/google/flax/pull/2856

Full Changelog: https://github.com/google/flax/compare/v0.6.4...v0.6.5

What's Changed

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

New Contributors

• @wookayin made their first contribution in https://github.com/google/flax/pull/2697
• @vwxyzjn made their first contribution in https://github.com/google/flax/pull/2780
• @arcAman07 made their first contribution in https://github.com/google/flax/pull/2796
• @nouiz made their first contribution in https://github.com/google/flax/pull/2810
• @mathisgerdes made their first contribution in https://github.com/google/flax/pull/2578

Full Changelog: https://github.com/google/flax/compare/v0.6.3...v0.6.4