0.11.2

[cgarciae]

What's Changed

nnx.merge now doesn't create a copy of the Variables in the incoming states by default, meaning that the new merged structures holds references to the incoming Variables. This enables new patterns, for example its now possible to create models with the same state but with different runtime behavior:

model = SomeModel(...)
# create eval model
eval_model = nnx.merge(*nnx.split(model))  # same Variables, different structure
eval_model.eval()

model and eval_model share the same Variables and are therefore kept in sync but have different runtime behavior, this avoids having to constantly mutate a single model back and forth between different runtime modes which can be error prone / cause unwanted recompilation.

To keep the old behavior use nnx.merge(..., copy=True).

PRs

• add Rngs random helpers by @cgarciae in add Rngs random helpers #4876
• Fix re-export and docs for identity by @jlperla in Fix re-export and docs for identity #4850
• Fix ToLinen docstring return description by @mohsinm-dev in Fix ToLinen docstring return description #4852
• Update doc build instructions and clean up unused packages by @IvyZX in Update doc build instructions and clean up unused packages #4885
• Improve docs related with dataclasses by @IvyZX in Improve docs related with dataclasses #4884
• Fix broken contributing documentation link by @mohsinm-dev in Fix broken contributing documentation link #4855
• Internal change by @copybara-service[bot] in Internal change #4886
• Fix string key preservation in replace_by_pure_dict by @mohsinm-dev in Fix string key preservation in replace_by_pure_dict #4860
• Remove the need for Conv and ConvTranspose to know the precise batch size. by @copybara-service[bot] in Remove the need for Conv and ConvTranspose to know the precise batch size. #4877
• call jax's source_info_util.register_exclusion in flax's traceback_util.register_exclusion by @copybara-service[bot] in call jax's source_info_util.register_exclusion in flax's traceback_util.register_exclusion #4887
• Update typo in nnx.Optimizer by @codinfox in Update typo in nnx.Optimizer #4880
• Exposed split_rngs docstring in the docs_nnx by @vfdev-5 in Exposed split_rngs docstring in the docs_nnx #4846
• Pin sentencepiece version to 0.2.0 to fix head by @IvyZX in Pin sentencepiece version to 0.2.0 to fix head #4892
• Relax duplicate check to exclude non-string values such as PartitionSpec.UNCONSTRAINED, since those can be repeated. by @copybara-service[bot] in Relax duplicate check to exclude non-string values such as PartitionSpec.UNCONSTRAINED, since those can be repeated. #4881
• add find_duplicates by @cgarciae in add find_duplicates #4894
• Sharding API improvements (non breaking) by @IvyZX in Sharding API improvements (non breaking) #4893
• document jax.random shorthand methods by @cgarciae in document jax.random shorthand methods #4899
• Optimiser was already instantiated using the model - 05_vae.py by @nenuadrian in Optimiser was already instantiated using the model - 05_vae.py #4857
• revert is_leaf logic in _check_carry_same_references by @copybara-service[bot] in revert is_leaf logic in _check_carry_same_references #4903
• Doc fix: remove outdated advice on flax v0.6.10; it was released two years ago. by @copybara-service[bot] in Doc fix: remove outdated advice on flax v0.6.10; it was released two years ago. #4910
• Fix bug when raising ScopeParamNotFoundError. by @copybara-service[bot] in Fix bug when raising ScopeParamNotFoundError. #4898
• fix mypy on main by @cgarciae in fix mypy on main #4909
• merge no copy Variables by @cgarciae in merge no copy Variables #4912
• update version to 0.11.2 by @copybara-service[bot] in update version to 0.11.2 #4915

New Contributors

• @mohsinm-dev made their first contribution in Fix ToLinen docstring return description #4852
• @codinfox made their first contribution in Update typo in nnx.Optimizer #4880
• @nenuadrian made their first contribution in Optimiser was already instantiated using the model - 05_vae.py #4857

Full Changelog: v0.11.1...v0.11.2

---

This discussion was created from the release 0.11.2.

[HeavyCrab]

Will this change affect nnx.update(model, new_model_state)?