Version 0.10.2

What's Changed

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

New Contributors

• @jlperla made their first contribution in #4351

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