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

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