Version 0.10.3

[cgarciae]

What's Changed

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

New Contributors

• @tilakrayal made their first contribution in Fixed the broken link in haiku_to_flax.rst file #4402
• @emmanuel-ferdman made their first contribution in Update async_checkpointer.py reference #4385
• @Matt-Hurd made their first contribution in Fix multiple links to Orbax documentation #4364

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

---

This discussion was created from the release Version 0.10.3.