v0.11.0

v0.11.0 - Pytrees, MutableArrays, and more!

This version of Flax introduces some changes to improve interop with native JAX and adds support for the new jax.experimental.MutableArray. More on this soon! However, some breaking changes to align with the JAX way of doing things were necessary. Most code should remain intact, however, the following changes deviate from the current behavior:

• Rngs in standard layers: all standard layers no longer hold a shared reference to the rngs object given in the constructor, instead they now keep a fork-ed copy of the Rngs or RngStream objects. This impacts Using Rngs in NNX Transforms and Loading Checkpoints with RNGs.
• Optimizer Updates: the Optimizer abstraction no longer holds a reference to the model to avoid reference sharing, instead the model must be provided as the first argument to update.
• Modules as Pytrees: Modules are now pytrees! This avoid unnecessary use of split and merge when interacting trivially with raw JAX transforms (state must still be manually propagated if not using MutableArrays, and referential transparency is still an issue). This affects when operating on Pytrees containing NNX Objects with jax.tree.* APIs.

Checkout the full NNX 0.10 to NNX 0.11 migration guide.

In the near future we'll share more information about new ways of using NNX with JAX transforms directly by leveraging the new Pytree and MutableArray support. Stay tuned!

What's Changed

• [nnx] mutable array p3 by @cgarciae in #4755
• [nnx] allow method calls in ToLinen by @cgarciae in #4808
• Internal change by @copybara-service[bot] in #4807
• Preserve sharding information in axes_scan by @copybara-service[bot] in #4806
• Deduplicate contributing and philosophy and move to main site by @IvyZX in #4809
• Fixed nnx.remat docstring rendering by @vfdev-5 in #4790
• Added a note to gemma guide about model's license consent on kaggle by @vfdev-5 in #4776
• [nnx] ToLinen add abtract_init flag by @cgarciae in #4813
• Modify NNX to use id(variable) instead of nnx.Variables as dictionary by @divyashreepathihalli in #4814
• Allow using LazyRngs for flax init/apply. by @copybara-service[bot] in #4818
• [nnx] remove VariableState by @cgarciae in #4800
• Fix failing CI jobs: trailing whitespace, deprecated .type usage by @vfdev-5 in #4823
• [nnx] fix Rngs dtype check by @cgarciae in #4820
• refactor: move usages of .value to [...] in modules_test.py by @lukeyeh in #4815
• Added training script for Gemma model by @vfdev-5 in #4822
• [nnx] add flax_pytree_module flag by @cgarciae in #4811
• create ModelAndOptimizer symbol by @copybara-service[bot] in #4849
• [nnx] remove Optimizer.model attribute by @cgarciae in #4842
• [nnx] add mutable array support in update by @cgarciae in #4851
• Migrate transforms_test.py from .value to [...] by @lukeyeh in #4841
• 0.11.0 migration guide by @cgarciae in #4854

New Contributors

• @divyashreepathihalli made their first contribution in #4814
• @lukeyeh made their first contribution in #4815

Full Changelog: v0.10.7...v0.11.0

0.10.7

What's Changed

• Added identity export from JAX by @jlperla in #4652
• Fixes a bug in type annotations for scope.param (unbox=True should accept callable[..., T | AxisMEtadata[T]] and return T, while unbox=False should always return the same thing as what callable returning. by @copybara-service in #4727
• fix merge by @copybara-service in #4731
• [nnx] make Variable a pytree by @cgarciae in #4728
• [nnx] add JitWrapped API by @cgarciae in #4699
• Update JAX nightly index usage by @copybara-service in #4733
• [nnx] mutable array p1 by @cgarciae in #4715
• add dataclass by @copybara-service in #4739
• [flax] unconditionally register nnx.Variable as a pytree by @copybara-service in #4748
• Updated version of pre-commit-hooks in .pre-commit-config.yaml by @vfdev-5 in #4746
• Fixed docstring visibility for nnx.eval_shape by @vfdev-5 in #4747
• Added keep_rngs arg to MHA to optionally store rngs by @vfdev-5 in #4749
• MultiHeadAttention only keeps rngs if dropout_rate is positive by @copybara-service in #4750
• [nnx] mutable array p2 by @cgarciae in #4741
• Add in_kv_features argument to nnx.MultiHeadAttention, addressing #4756. by @copybara-service in #4757
• Fix broken link for Transforms guide by @nireekshak in #4763
• Minor improvements of lm1b_nnx example by @vfdev-5 in #4745
• Fix head CI tests by @IvyZX in #4764
• Fix typos by @nireekshak in #4725
• Check for leaves of type variablelib.Variable when getting sharding specs. by @copybara-service in #4769
• Fixes #1925 non-str dict keys not suppoted in module state by @muhrin in #4563
• Modified the Functional API link by @nireekshak in #4767
• Fix hardcoded link to filter guide in docs by @hamogu in #4768
• Fix bad doc links by @IvyZX in #4770
• revise axes_scan to flatten argument pytrees only once by @copybara-service in #4772
• Simplify ToNNX access of Linen module methods by @IvyZX in #4766
• Use .input_formats and .output_formats in place of .input_layouts and .output_layouts respectively. by @copybara-service in #4784
• Exposed OptState in nnx module by @vfdev-5 in #4788
• Fixes colab link for nnx docs by @vfdev-5 in #4775
• Internal changes by @copybara-service in #4786
• Fix typo in Flax nnx_basics doc. by @copybara-service in #4781
• update version to 0.10.7 by @cgarciae in #4798

New Contributors

• @nireekshak made their first contribution in #4763
• @muhrin made their first contribution in #4563
• @hamogu made their first contribution in #4768

Full Changelog: v0.10.6...v0.10.7

0.10.6

What's Changed

• Sow top activations based on absolute value. by @copybara-service in #4670
• Add support for layer-specific rope scale factors. by @copybara-service in #4672
• Automatic model selection for Gemma 3 models. by @copybara-service in #4671
• Make LoRA's dtype arg useful by @IvyZX in #4681
• [NVIDIA] Support FP8 Einsum Op by @kaixih in #4686
• [nnx] remove deprecated APIs by @cgarciae in #4627
• Add attention_bias parameter to MultiHeadDotProductAttention. by @copybara-service in #4694
• Unit tests for attention_bias parameter to MultiHeadDotProductAttention. Add parameter to all overloads to make pytype happy. by @copybara-service in #4702
• Rollback of attention_bias parameter, because the change overrides the attention bias for injected attention functions. by @copybara-service in #4703
• Add custom einsum op to Einsum() by @IvyZX in #4705
• [nnx] refactor GraphDef by @cgarciae in #4630
• Make fully replicated array before saving checkpoints for examples that use pmap. by @copybara-service in #4707
• Fix CI by @cgarciae in #4716
• remove "nnx" collection in ToLinen by @copybara-service in #4708
• [nnx] flaxlib types by @cgarciae in #4639
• v0.10.6 by @cgarciae in #4724

Full Changelog: v0.10.5...v0.10.6

0.10.5

What's Changed

• [nnx] fix tabulate by @cgarciae in #4580
• Refactor bridge.Module tests from wrappers_test.py to another file. by @copybara-service in #4581
• Avoid calls to jnp.shape for non-array inputs. by @jakevdp in #4592
• remove Embed nan casting by @cgarciae in #4600
• Add QK Norm. by @copybara-service in #4594
• Util to let bridge module work with NNX submodules by @IvyZX in #4584
• Add configurable Query Pre Attention scalar. by @copybara-service in #4595
• Make RoPE Base Frequency configurable. by @copybara-service in #4596
• [nnx] pytrees are graph nodes by @cgarciae in #4547
• Add option to load checkpoints with transposed Gating Einsum. by @copybara-service in #4597
• add top_p sampling in gemma example by @copybara-service in #4591
• Fix position and name of Post Attention Norm. by @copybara-service in #4598
• Add Sow Config to from_params constructor. by @copybara-service in #4599
• bridge module with linen submodule by @IvyZX in #4604
• Dramatically speed up sampling compilation time by @copybara-service in #4574
• [nnx] improve grad docs by @cgarciae in #4588
• [nnx] add support for standalone Variables by @cgarciae in #4606
• add promote_dtype as a config option for multiple layers by @cgarciae in #4613
• Copybara import of the project: by @copybara-service in #4616
• Fixed typo in beam_search loop. by @copybara-service in #4615
• support swap model params in gemma sampler by @copybara-service in #4614
• Allow bridge module to have 'name' field by @IvyZX in #4619
• fix performance guide by @cgarciae in #4621
• Copybara import of the project: by @copybara-service in #4618
• Add REFLECT padding to convolution layer by @sarlinpe in #4553
• fix trace-level detection by @cgarciae in #4527
• Add attribute path customization to bridge modules by @IvyZX in #4624
• add reprlib max depth flag by @cgarciae in #4632
• Allow custom axis metadata annotation during transforms by @IvyZX in #4637
• [bridge module] Allow name arg to represent actual submodule path by @IvyZX in #4634
• [nnx] improve Variable proxy for binary operations by @cgarciae in #4641
• Fix module stack typing annotation. by @copybara-service in #4633
• Stop passing reduce_axes to jax.grad, jax.vjp, and jax.value_and_grad. by @copybara-service in #4617
• discord release webhook by @cgarciae in #4646
• [nnx] support Array leaves in graph nodes by @cgarciae in #4612
• Roll up package jax version and uv.lock by @IvyZX in #4648
• Use jax.nn.dot_product_attention when possible by @IvyZX in #4649
• Fix flaky vmap test tolerance. by @copybara-service in #4653
• Test runner ubuntu upgrade 24.04 by @IvyZX in #4659
• Fix lazy_init typo by @IvyZX in #4657
• deflake a test by @copybara-service in #4663
• v0.10.5 by @cgarciae in #4656

New Contributors

• @sarlinpe made their first contribution in #4553

Full Changelog: v0.10.4...v0.10.5

Release 0.10.4

What's Changed

• update pypi publish by @cgarciae in #4538
• [nnx] register_variable_name refactor by @copybara-service in #4540
• added support to the accuracy metric for binary classification by @mattbahr in #4536
• [nnx] bridge Module by @cgarciae in #4542
• [nnx] copy _var_metadata by @copybara-service in #4548
• [bridge] fix unbox logic by @copybara-service in #4551
• Add is_initializing API by @copybara-service in #4550
• [nnx] Add specific model typing for nnx.Optimizer by @marcelroed in #4470
• Add linen metadata conversion to linx by @IvyZX in #4552
• [bridge] improve Module context by @cgarciae in #4554
• Raise error if user uses 'name' in bridge module setup by @IvyZX in #4555
• Add deprecation warning to all nnx.State methods by @IvyZX in #4561
• [nnx] add shard_map by @cgarciae in #4490
• Fix CI breakages from newest jax by @IvyZX in #4576
• [bridge] Set _initializing correctly and avoid return RNG states by @copybara-service in #4569
• v0.10.4 by @cgarciae in #4579

New Contributors

• @mattbahr made their first contribution in #4536
• @marcelroed made their first contribution in #4470

Full Changelog: v0.10.3...v0.10.4

Version 0.10.3

What's Changed

• Fix fori_loop and while_loop on multiple modules by @IvyZX in #4390
• Upgrade Flax readme to NNX by @8bitmp3 in #4386
• [nnx] add performance guide notebook by @cgarciae in #4384
• Automated Code Change by @copybara-service in #4393
• [nnx] optimize NodeDef.attributes by @cgarciae in #4399
• Fixed the broken link in haiku_to_flax.rst file by @tilakrayal in #4402
• [nnx] optimize Variable by @cgarciae in #4400
• Update Flax NNX Randomness by @8bitmp3 in #4279
• Remove the repeated methods in flax.nnx.Module documentation by @rajasekharporeddy in #4416
• Fixed the broken link in linen_to_nnx.rst by @tilakrayal in #4415
• [nnx] add FlatState by @cgarciae in #4410
• Update async_checkpointer.py reference by @emmanuel-ferdman in #4385
• Fix multiple links to Orbax documentation by @Matt-Hurd in #4364
• Update links in Why Flax NNX documentation by @rajasekharporeddy in #4425
• [nnx] fix transforms guide by @cgarciae in #4421
• Add benchmark on state traversal, and a readme by @IvyZX in #4428
• Update Flax NNX performance guide by @8bitmp3 in #4401
• Create sharding via Partitioned.get_sharding() by @copybara-service in #4427
• Update Flax NNX vs JAX Transformations guide by @8bitmp3 in #4286
• Upgrade Flax NNX Gemma Sampling Inference doc by @8bitmp3 in #4325
• Update NNX merge docs in graph.py by @8bitmp3 in #4411
• Fix main and add nnx.fori_loop test by @cgarciae in #4472
• Upgrade Flax NNX Filters doc by @8bitmp3 in #4199
• Makes flax Modules more compatible with IPython auto-reload. by @copybara-service in #4420
• [nnx] RNN: add broadcast_rngs and state_axes APIs by @cgarciae in #4407
• Allow nnx.bridge.variables.nnx_attrs_to_linen_vars take nnx.VariableState as argument. by @copybara-service in #4473
• [nnx] add state summaries for print and display by @cgarciae in #4438
• Copybara import of the project: by @copybara-service in #4475
• [nnx] add state summaries for print and display by @copybara-service in #4477
• CI: add scheduled test against nightly JAX releases by @jakevdp in #4478
• CI: pin actions to specific commits by @jakevdp in #4479
• [nnx] fix MultiMetric typing by @cgarciae in #4485
• [nnx] fix ToNNX linen_attributes update by @cgarciae in #4486
• Remove usages of orbax_utils.save_args_from_target, as this function does nothing (it used to control a checkpointing behavior that has since been optimized away). by @copybara-service in #4482
• [nnx] improve Module docs by @cgarciae in #4499
• Update einsum layer for Gemma example by @copybara-service in #4498
• [nnx] fix fiddle by @cgarciae in #4500
• Don't create param in normalization layers instead of create None-value params. by @copybara-service in #4501
• Rename variable string mapping utils and move them to variableslib by @IvyZX in #4503
• fix LoRA initialization error in nnx layer by @copybara-service in #4502
• Remove all Param(None) lines from NNX by @IvyZX in #4504
• make gemma FFW LoRA friendly by @copybara-service in #4510
• Add nnx.Module.perturb by @IvyZX in #4515
• [nnx] add tabulate by @cgarciae in #4493
• batch_norm.rst: == should be = by @cool-RR in #4524
• v0.10.3 by @cgarciae in #4525

New Contributors

• @tilakrayal made their first contribution in #4402
• @emmanuel-ferdman made their first contribution in #4385
• @Matt-Hurd made their first contribution in #4364

Full Changelog: v0.10.2...v0.10.3

Version 0.10.2

What's Changed

• Add nnx.fori_loop by @IvyZX in #4353
• Linesearch (and lbfgs) support by @jlperla in #4351
• Upgrade Flax NNX Haiku Linen migration doc by @8bitmp3 in #4200
• Fix PRNG handling in nn.jit under nn.scan. by @copybara-service in #4359
• support passing arguments directly to the struct.dataclass decorator by @copybara-service in #4275
• Avoid assert_array_equal for PRNG keys. by @copybara-service in #4363
• [nnx] support pure dicts by @cgarciae in #4352
• [nnx] add data parallel toy example by @cgarciae in #4354
• Add logical axis global context support for NNX by @IvyZX in #4350
• [nnx] fix ToLinen kwargs by @copybara-service in #4270
• [nnx] use HashableMapping instead of FrozenDict by @cgarciae in #4376
• [nnx] fix while_loop/fori_loop bug when sharing references by @cgarciae in #4379
• Add flax.nnx.eval_shape docstring by @8bitmp3 in #4374
• Setup the flaxlib in C++, using Meson and Nanobind. by @copybara-service in #4380
• Add flax.nnx.remat docstring by @8bitmp3 in #4373
• [nnx] add checkify by @cgarciae in #4381
• Lint flax.nnx.while_loop docstring by @8bitmp3 in #4371
• Lint flax.nnx.fori_loop docstring by @8bitmp3 in #4370
• [nnx] add some optimizations to graph.py by @cgarciae in #4377
• update version to 0.10.2 by @cgarciae in #4387

New Contributors

• @jlperla made their first contribution in #4351

Full Changelog: v0.10.1...v0.10.2

Version 0.10.1

What's Changed

• Add Flax NNX GraphDef docstring by @8bitmp3 in #4302
• Flesh out the Haiku/Flax guide by @IvyZX in #4305
• [nnx] improve mnist tutorial by @cgarciae in #4316
• Update Flax Evolution from Linen to NNX guide by @8bitmp3 in #4289
• [nnx] try casting integers keys in State.replace_by_pure_dict by @cgarciae in #4317
• Fixed nnx examples bad links in the README.md by @vfdev-5 in #4282
• Fix philosophy link by @jorisSchaller in #4313
• [nnx] add gemma notebook by @cgarciae in #4075
• [nnx] improve init_cache docs by @cgarciae in #4291
• remove markdown from section titles by @cgarciae in #4322
• Avoid depending on JAX internals, which are about to change. by @copybara-service in #4326
• Remove outdated compatibility code. by @jakevdp in #4324
• fix ruff complaints by @levskaya in #4331
• Remove GeGLU activation function and golden tests. by @copybara-service in #4303
• Avoid using float32 in normalization for mean/var and scale/bias parameters when force_float32_reductions=False by @copybara-service in #4314
• Avoid assert_array_equal on PRNG keys. by @jakevdp in #4332
• Fix typos in Flax NNX Migrating from Haiku to Flax by @8bitmp3 in #4337
• Add API reference for flax.nnx.nn and improve landing page by @IvyZX in #4338
• [nnx] improve transforms guide by @cgarciae in #4333
• [nnx] cleanup gemma notebook by @cgarciae in #4334
• Remove non-lazy RNG compat mode and flag from flax. by @copybara-service in #4339
• [nnx] fix custom_vjp by @cgarciae in #4306
• Define model surgery in docs by @8bitmp3 in #4349
• [nnx] update State and variables docstrings by @cgarciae in #4346
• Add NNX transforms nnx.while_loop and nnx.switch by @IvyZX in #4343
• update version to v0.10.1 by @cgarciae in #4345

New Contributors

• @vfdev-5 made their first contribution in #4282

Full Changelog: v0.10.0...v0.10.1

Version 0.10.0

What's Changed

• [nnx] clear nnx basics pip logs by @cgarciae in #4149
• Support linen <-> nnx metadata box converging in nnx.bridge by @IvyZX in #4145
• Add nnx bridge API reference to site by @IvyZX in #4158
• [nnx] use jax-style transforms API in nnx_basics by @cgarciae in #4155
• [nnx] improve nnx.scan in_axes/out_axes by @cgarciae in #4157
• Support direct quantization for FP8 matmul by @wenscarl in #3922
• Upgrade Flax NNX Model Surgery by @8bitmp3 in #4135
• [nnx] add more Variable proxy methods by @cgarciae in #4170
• [nnx] disallow Array leaves by @copybara-service in #4172
• Internal change by @copybara-service in #4176
• [nnx] improve landing page and nnx_basics messaging by @cgarciae in #4168
• Fixes a small bug in flax.linen.share_scope, where the scopes of children of the module being merged that were created before setup(),were not being updated to point to the new scope, and so they would end up staying under the original tree. by @copybara-service in #4150
• Move all NNX content up a level to be equal with Linen, to make python packaging more consistent. by @copybara-service in #4177
• Add a guide for nnx.bridge by @IvyZX in #4171
• [nnx] improve Optimizer metadata propagation by @cgarciae in #4180
• [nnx] enable sharding transformation on integer prefixes by @cgarciae in #4185
• Support linen.LogicallyPartitioned <-> nnx.Variable by @IvyZX in #4161
• Clean up axis hooks in nnx.Variable by @IvyZX in #4189
• Merge nnx.errors to flax.errors by @IvyZX in #4186
• [nnx] optimize jit by @cgarciae in #4191
• Split documentation for Linen and NNX by @cgarciae in #4192
• Partially revert #4192 which sets back a bunch of previous merged pushes. by @copybara-service in #4201
• Align bridge variable tree structures by @IvyZX in #4194
• [NNX site] Fix landing page and banner phrasing and add examples page by @IvyZX in #4202
• shorten banners by @cgarciae in #4206
• Add trimmed Linen to NNX guide by @IvyZX in #4209
• Minor documentation fixes for AxisMetadata. by @copybara-service in #4178
• fix tests for numpy 2.0 compatibility by @copybara-service in #4215
• Forward all arguments when using nnx.transforms.deprecated.scan as a decorator. by @copybara-service in #4208
• [nnx] add transforms guide by @cgarciae in #4197
• [nnx] fix transforms guide by @cgarciae in #4223
• Flax NNX GSPMD guide by @IvyZX in #4220
• Update libraries to use JAX's limited (and ill-advised) trace-state-querying APIs rather than depending on JAX's deeper internals, which are about to change. by @copybara-service in #4225
• [nnx] add Randomness guide by @cgarciae in #4216
• Add pure dict conversion util functions to nnx.State. by @IvyZX in #4230
• [nnx] Simplify traversal by @cgarciae in #4205
• Fix false positive tracer leaks in flax library. by @copybara-service in #4232
• [nnx] add flaxlib by @copybara-service in #4235
• [nnx] improve docs by @cgarciae in #4236
• point nnx banner to flax-linen by @cgarciae in #4237
• update banners by @cgarciae in #4238
• Fix scale dtype and refactor q_dot_dq by @wenscarl in #4229
• update banners by @cgarciae in #4241
• Add redirects for Linen guide links in the NNX site scope. by @IvyZX in #4242
• Internal change by @copybara-service in #4243
• Copybara import of the project: by @copybara-service in #4245
• Update Flax NNX Scale Up SPMD guide by @8bitmp3 in #4239
• Upgrade Flax NNX basics doc by @8bitmp3 in #4173
• Improve landing page, glossary and misc by @IvyZX in #4244
• Nitting and adding links by @8bitmp3 in #4248
• enable doctest on notebooks by @cgarciae in #4250
• Update index.rst by @ariG23498 in #4251
• Add NNX checkpointing guide by @IvyZX in #4249
• Add checkpointing guide to website index. by @copybara-service in #4263
• Update to Flax NNX Transforms doc by @8bitmp3 in #4264
• Add why nnx by @cgarciae in #4240
• [nnx] add cloudpickle support by @cgarciae in #4253
• Fix typo: impost to import by @Vilin97 in #4256
• [nnx] revive TrainState toy example by @cgarciae in #4226
• [nnx] add custom_vjp to docs by @cgarciae in #4266
• remove flax-nnx urls by @cgarciae in #4267
• Add flatten to nnx.graph autosummary in graph.rst by @8bitmp3 in #4255
• [nnx] add FSDP toy example with custom optimizer by @cgarciae in #4183
• Update Flax NNX Landing Page by @8bitmp3 in #4274
• Update to Flax NNX Model Surgery by @8bitmp3 in #4276
• Update Why Flax NNX guide by @8bitmp3 in #4262
• Update to Flax NNX MNIST tutorial by @8bitmp3 in #4277
• [nnx] improve randomness guide by @cgarciae in #4281
• Remove notebook exceptions in docs_nnx doctest by @IvyZX in #4285
• [nnx] add PrefixMapping by @cgarciae in #4278
• [nnx] state filters by @cgarciae in #4288
• Fix devcontainer setup by @jorisSchaller in #4299
• Ugrade Flax NNX Checkpointing guide by @8bitmp3 in #4294
• Update Flax NNX Scale Up guide by @8bitmp3 in #4296
• Porting RNN from Linen to NNX by @zinccat in #4272
• Update Flax NNX Glossary by @8bitmp3 in #4284
• update version to 0.10.0 by @cgarciae in #4292

New Contributors

• @ariG23498 made their first contribution in #4251
• @Vilin97 made their first contribution in #4256
• @jorisSchaller made their first contribution in #4299
• @zinccat made their first contribution in #4272

Full Changelog: v0.9.0...v0.10.0

v0.9.0

What's Changed

• Add NNX surgery guide by @IvyZX in #4005
• Port gemma/transformer to NNX by @copybara-service in #4019
• upgrade python to 3.10 + use pyupgrade by @cgarciae in #4038
• [nnx] add Using Filters guide by @cgarciae in #4028
• v0.8.6 by @cgarciae in #4040
• allow imagenet training profiling to be disabled in config by @copybara-service in #4043
• [nnx] LoRAParam inherits from Param by @cgarciae in #3988
• [linen] allows multiple compact methods by @cgarciae in #3808
• Added support of NANOO fp8. by @wenchenvincent in #3993
• Add functool.wraps() annotation to flax.nn.jit. by @copybara-service in #4051
• Fix typo in nnx_basics doc by @rajasekharporeddy in #4047
• [nnx] fix Variable overloads and add shape/dtype properties by @cgarciae in #4049
• Stop writing msgpack file for new checkpoints and update empty nodes handling so that it no longer depends on this file. by @copybara-service in #4039
• [nnx] stabilize unsafe_pytree by @cgarciae in #4030
• Stop writing msgpack file for new checkpoints and update empty nodes handling so that it no longer depends on this file. by @copybara-service in #4055
• [NVIDIA] Rename fp8 custom dtype to fp32_max_grad by @kaixih in #3984
• [nnx] fix mnist_tutorial colab link by @cgarciae in #4063
• [nnx] fix Accuracy on eager mode by @cgarciae in #4065
• Update orbax_upgrade_guide.rst for async checkpointing usage examples by @kaushaladiti-2802 in #4036
• Re-enable some tests after Python 3.9 is dropped by @IvyZX in #4067
• Rename nnx.compat to nnx.bridge by @IvyZX in #4066
• [nnx] improve mnist tutorial by @cgarciae in #4070
• Modify Flax checkpointing in preparation for cl/650338576. by @copybara-service in #4072
• Remove some outdated backward-compatibility code. by @copybara-service in #4068
• [NVIDIA] Add a user guide for fp8 by @kaixih in #4076
• [nnx] add extract APIs by @cgarciae in #4078
• [example]: remove lm1b useless parallism rules by @knightXun in #4077
• [nnx] improve filters guide by @cgarciae in #4059
• [nnx] add call by @cgarciae in #4004
• Ignore Orbax warning in deprecated flax.training.checkpoints.py to unbreak head doctest by @IvyZX in #4092
• fix mypy failures due tu numpy update by @cgarciae in #4098
• [linen] generalize transform caching by @copybara-service in #4057
• [linen] fold rngs on jit to improve caching by @copybara-service in #4064
• Add shape-based lazy init to LinenToNNX (prev LinenWrapper) by @IvyZX in #4081
• [nnx] add reseed by @cgarciae in #4099
• [nnx] add split/merge_inputs by @cgarciae in #4084
• Perform shape checks for self.param AFTER unboxing by @danielwatson6 in #4079
• fix restore_checkpoint example in docstring by @copybara-service in #4101
• [numpy] Fix users of NumPy APIs that are removed in NumPy 2.0. by @copybara-service in #4104
• set profile_duration_ms = None as in periodic_actions there's default value for both num_profile_steps and profile_duration_ms, and the profile stopping condition is when both num_profile_steps and profile_duration_ms are satisfied, so setting profile_duration_ms=None so that the passed num_profile_steps value gets used by @copybara-service in #4096
• [linen] add share_scope by @cgarciae in #4102
• Allow metadata pass-through in flax.struct.field by @cool-RR in #4056
• avoid mixing einsum_dot_general and einsum argument by specifying them explicitly in the caller. by @copybara-service in #4115
• Add logging to track deprecated codepaths. by @copybara-service in #4121
• [pmap no rank reduce cleanup]: When flipping the by @copybara-service in #4125
• Add NNXToLinen wrapper to nnx.bridge by @IvyZX in #4126
• Switch NNX to use Treescope instead of Penzai. by @copybara-service in #4132
• Add GroupNorm to NNX normalization layers by @treigerm in #4095
• [nnx] fix initializing propagation by @cgarciae in #4134
• add JAX-style NNX Transforms FLIP by @cgarciae in #4108
• Fix _ParentType annotation by @dcharatan in #4120
• add uv.lock file by @copybara-service in #4139
• use uv package manager by @cgarciae in #4136
• More testing and misc fixes on wrappers by @IvyZX in #4137
• Fix link to orbax documentation by @cool-RR in #4123
• [nnx] experimental transforms by @cgarciae in #3963
• [nnx] improve docs by @cgarciae in #4141
• remove repeated license headers by @cgarciae in #4148
• update Flax to version 0.9.0 by @copybara-service in #4147

New Contributors

• @wenchenvincent made their first contribution in #3993
• @rajasekharporeddy made their first contribution in #4047
• @kaushaladiti-2802 made their first contribution in #4036
• @knightXun made their first contribution in #4077
• @danielwatson6 made their first contribution in #4079
• @cool-RR made their first contribution in #4056
• @treigerm made their first contribution in #4095
• @dcharatan made their first contribution in #4120

Full Changelog: v0.8.5...v0.9.0