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