0.12.5

What's Changed

• maintain data/static definition in split / state by @copybara-service[bot] in #5243
• Add graph=False support for nnx.grad and nnx.value_and_grad by @copybara-service[bot] in #5240
• Add tree-mode support for nnx.remat by @copybara-service[bot] in #5242
• Add tree-mode support for nnx.vmap by @copybara-service[bot] in #5250
• Add tree-mode arg error handling by @copybara-service[bot] in #5251
• [flax:benchmarks] Fix flax benchmarks. by @copybara-service[bot] in #5241
• Improve unregistered data detection by @copybara-service[bot] in #5249
• Add graph argument to split, state, graphdef, clone by @copybara-service[bot] in #5258
• Add nnx_graph_mode config flag by @copybara-service[bot] in #5261
• Add tree-mode support to nnx.{cond,switch,eval_shape,checkify} by @copybara-service[bot] in #5252
• Add tree-mode support to nnx.scan by @copybara-service[bot] in #5262
• Add tree-mode support for nnx.custom_vjp by @copybara-service[bot] in #5264
• Add tree-mode support to nnx.{while,fori}_loop by @copybara-service[bot] in #5263
• add tree-mode support for pmap by @copybara-service[bot] in #5265
• Add tree-mode support to iter_graph by @copybara-service[bot] in #5266
• Add tree-mode support for nnx.recursive_map, nnx.view, and nnx.view_info by @copybara-service[bot] in #5267
• Use linen_vars_to_nnx_attrs in ToLinen variable restoration by @copybara-service[bot] in #5272
• Add tree-mode-only nnx.{vjp,jvp,jit_partial} by @copybara-service[bot] in #5268
• Use linen_vars_to_nnx_attrs in ToLinen variable restoration by @copybara-service[bot] in #5273
• fix nn.vmap's _broadcast_prefix_tree by @copybara-service[bot] in #5282
• fix nn.vmap's split_fn by @copybara-service[bot] in #5286
• [pmap] Remove jax.config.pmap_shmap_merge. by @copybara-service[bot] in #5289
• add set_graph_mode by @copybara-service[bot] in #5269
• add error handling to iter_children by @copybara-service[bot] in #5296
• expose and update view guide by @cgarciae in #5294
• Remove chex dependency by @samanklesaria in #5295
• improve error messages for tree mode errors by @copybara-service[bot] in #5301
• remove experimental Variable.mutable property by @copybara-service[bot] in #5304
• update version to 0.12.5 by @copybara-service[bot] in #5305
• simplify nnx.display for Variable by @copybara-service[bot] in #5306
• simplify set_view API by @copybara-service[bot] in #5303

Full Changelog: v0.12.4...v0.12.5