Flax Files

What's Changed

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

Full Changelog: https://github.com/google/flax/compare/v0.12.4...v0.12.5