Version 0.10.0

[IvyZX]

What's Changed

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

New Contributors

• @ariG23498 made their first contribution in Update index.rst #4251
• @Vilin97 made their first contribution in Fix typo: impost to import #4256
• @jorisSchaller made their first contribution in Fix devcontainer setup #4299
• @zinccat made their first contribution in Porting RNN from Linen to NNX #4272

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

---

This discussion was created from the release Version 0.10.0.