google/flax v0.8.0 on GitHub

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 #3571
@jianyizh made their first contribution in #3590

Full Changelog: v0.7.5...v0.8.0