Flax Versions Save

Flax is a neural network library for JAX that is designed for flexibility.

v0.8.3

2 weeks ago

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

v0.8.2

2 months ago

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

v0.8.1

3 months ago

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

v0.8.0

3 months ago

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

v0.7.5

6 months ago

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

v0.7.4

8 months ago

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

v0.7.3

8 months ago

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

v0.7.2

9 months ago

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

v0.7.1

9 months ago

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

v0.7.0

10 months ago

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