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