Version 0.5.3

What's Changed

• Adds .pre-commit-config.yaml by @copybara-service in #2212
• Fix missing passthrough of nn.scan unroll arg by @jheek in #2213
• Test Notebooks on CI by @cgarciae in #2166
• Bump numpy from 1.21.4 to 1.22.0 in examples by @marcvanzee in #2228
• Add nn.switch by @cgarciae in #2205
• Fix notebooks by @cgarciae in #2231
• Add launch section with colab button by @cgarciae in #2235
• Enabling the dollarmath extension of MyST to render correctly math expresions by @WaterKnight1998 in #2238
• Update codediff to use sphinx-design tabs by @cgarciae in #2204
• Fix tests by @cgarciae in #2253
• Add single-host async save to save_checkpoint. by @IvyZX in #2233
• Add a method for detecting the use of "init" functions. by @levskaya in #2234
• Small fix in MNIST example by @marcvanzee in #2258
• Fix typos in the doc of flax.linen.Module.bind by @nalzok in #2269
• Add colab button to flax_basics by @cgarciae in #2276
• Fix type annotations by @cgarciae in #2281
• Exclude pseudo-fields of dataclass by @YouJiacheng in #2199
• Fix variable aliasing in put_variable by @jheek in #2296
• Update reference to tree_map to avoid deprecation warning. by @copybara-service in #2298
• Fix nondeterministic bug arising from sharing logic during module adoption. by @copybara-service in #2302
• fix ppo example typo by @fuyw in #2306
• Forward axis_size tot jax.vmap by @jheek in #2310
• cleanup: replace deprecated jax.tree_map with jax.tree_util.tree_map by @copybara-service in #2311
• Add GlobalDeviceArray/multihost checkpoint support to Flax. by @copybara-service in #2287
• 0.5.3 update version & changelog by @IvyZX in #2330
• Replace use of id() with global counter-based id. by @levskaya in #2313

New Contributors

• @WaterKnight1998 made their first contribution in #2238
• @nalzok made their first contribution in #2269
• @YouJiacheng made their first contribution in #2199
• @fuyw made their first contribution in #2306

Full Changelog: v0.5.2...v0.5.3