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