Version 0.8.1

What's Changed

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

Full Changelog: v0.8.0...v0.8.1

v0.8.0

What's Changed

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

New Contributors

• @superbobry made their first contribution in #3448
• @VictorPrins made their first contribution in #3450
• @zhaoyang-0204 made their first contribution in #3462
• @Micky774 made their first contribution in #3498
• @HMUNACHI made their first contribution in #3512
• @carlosgmartin made their first contribution in #3510
• @giovannic made their first contribution in https://gith...

v0.7.5

What's Changed

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

New Contributors

• @wenscarl made their first contribution in #3322
• @kaixih made their first contribution in #3435

Full Changelog: v0.7.4...v0.7.5

v0.7.4

What's Changed

Added python version constraint >=3.9.

• Bump version no by 1 post-release. by @copybara-service in #3328
• Update minimum python version info. by @levskaya in #3331

Full Changelog: v0.7.3...v0.7.4

Version 0.7.3

What's Changed

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

New Contributors

• @peterdavidfagan made their first contribution in #3265
• @SauravMaheshkar made their first contribution in #3254

Full Changelog: v0.7.2...v0.7.3

Version 0.7.2

What's Changed

• make flax.core.copy add_or_replace optional by @PhilipVinc in #3241
• bump version to 0.7.2 by @chiamp in #3258
• Release 0.7.2 by @chiamp in #3269

Full Changelog: v0.7.1...v0.7.2

Version 0.7.1

What's Changed

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

New Contributors

• @marcoselvi made their first contribution in #3209

Full Changelog: v0.7.0...v0.7.1

Version 0.7.0

What's Changed

• delete long-unused dotgetter utility by @copybara-service in #3156
• Add sections to flax.linen toctree by @cgarciae in #3073
• update doctest requirements to use recentish jax minver by @levskaya in #3160
• Update gym to gymnasium by @cgarciae in #3133
• Add lifted transforms section to Haiku migration guide by @cgarciae in #3158
• relaxed tolerance by @chiamp in #3162
• fixed broken links by @chiamp in #3161
• added dict migration guide by @chiamp in #3109
• Update requirements, restructure files and fix formatting for VAE example by @canyon289 in #3046
• Update python version support by @cgarciae in #3168
• Improve Haiku migration guide by @cgarciae in #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 #3165
• fix precommit issues by @chiamp in #3170
• fixed incorrect reference link by @chiamp in #3167
• added absltest to linen_recurrent_test by @chiamp in #3172
• remove cell_size docs from RNN by @cgarciae in #3186
• 0.7.0 by @cgarciae in #3187

Full Changelog: v0.6.11...v0.7.0

Version 0.6.11

What's Changed

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

New Contributors

• @LenaMartens made their first contribution in #3154

Full Changelog: v0.6.10...v0.6.11

Version 0.6.10

What's Changed

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

New Contributors

• @Neilblaze made their first contribution in #3021

Full Changelog: v0.6.9...v0.6.10