Version 0.10.2

[IvyZX]

What's Changed

• Add nnx.fori_loop by @IvyZX in Add nnx.fori_loop #4353
• Linesearch (and lbfgs) support by @jlperla in Linesearch (and lbfgs) support #4351
• Upgrade Flax NNX Haiku Linen migration doc by @8bitmp3 in Upgrade Flax NNX Haiku Linen migration doc #4200
• Fix PRNG handling in nn.jit under nn.scan. by @copybara-service in Fix PRNG handling in nn.jit under nn.scan. #4359
• support passing arguments directly to the struct.dataclass decorator by @copybara-service in support passing arguments directly to the struct.dataclass decorator #4275
• Avoid assert_array_equal for PRNG keys. by @copybara-service in Avoid assert_array_equal for PRNG keys. #4363
• [nnx] support pure dicts by @cgarciae in [nnx] support pure dicts #4352
• [nnx] add data parallel toy example by @cgarciae in [nnx] add data parallel toy example #4354
• Add logical axis global context support for NNX by @IvyZX in Add logical axis global context support for NNX #4350
• [nnx] fix ToLinen kwargs by @copybara-service in [nnx] fix ToLinen kwargs #4270
• [nnx] use HashableMapping instead of FrozenDict by @cgarciae in [nnx] use HashableMapping instead of FrozenDict #4376
• [nnx] fix while_loop/fori_loop bug when sharing references by @cgarciae in [nnx] fix while_loop/fori_loop bug when sharing references #4379
• Add flax.nnx.eval_shape docstring by @8bitmp3 in Add flax.nnx.eval_shape docstring #4374
• Setup the flaxlib in C++, using Meson and Nanobind. by @copybara-service in Setup the flaxlib in C++, using Meson and Nanobind. #4380
• Add flax.nnx.remat docstring by @8bitmp3 in Add flax.nnx.remat docstring #4373
• [nnx] add checkify by @cgarciae in [nnx] add checkify #4381
• Lint flax.nnx.while_loop docstring by @8bitmp3 in Lint flax.nnx.while_loop docstring #4371
• Lint flax.nnx.fori_loop docstring by @8bitmp3 in Lint flax.nnx.fori_loop docstring #4370
• [nnx] add some optimizations to graph.py by @cgarciae in [nnx] add some optimizations to graph.py #4377
• update version to 0.10.2 by @cgarciae in update version to 0.10.2 #4387

New Contributors

• @jlperla made their first contribution in Linesearch (and lbfgs) support #4351

Full Changelog: v0.10.1...v0.10.2

---

This discussion was created from the release Version 0.10.2.