0.12.3

[cgarciae]

What's Changed

• Set numpy<2.4 to fix DeprecationWarning in the CI doctest by @vfdev-5 in Set numpy<2.4 to fix DeprecationWarning in the CI doctest #5163
• ignore optax deprecation warning by @cgarciae in ignore optax deprecation warning #5165
• fix general guides landing page by @ivrolan in fix general guides landing page #5139
• Remove nnx.split_rngs call wrapping nnx.scan in linen to nnx tutorial by @samanklesaria in Remove nnx.split_rngs call wrapping nnx.scan in linen to nnx tutorial #5160
• Add out_sharding arguments to linear layers where supported by @jackopenn in Add out_sharding arguments to linear layers where supported #5156
• fix kw_only_dataclasses for python 3.14 (part 2) by @copybara-service[bot] in fix kw_only_dataclasses for python 3.14 (part 2) #5135
• No public description by @copybara-service[bot] in No public description #5164
• Have _graph_flatten respect nnx.data declarations (extension of flatten respect nnx.data #5140) by @samanklesaria in Have _graph_flatten respect nnx.data declarations (extension of #5140) #5159
• [pmap] Avoid degraded performance under the new jax.pmap. by @copybara-service[bot] in [pmap] Avoid degraded performance under the new jax.pmap. #5152
• Support multiple None and UNCONSTRAINED when resolving logical rules by @copybara-service[bot] in Support multiple None and UNCONSTRAINED when resolving logical rules #5129
• improve hijax guide by @cgarciae in improve hijax guide #5115
• docs: fix typo 'paramater' -> 'parameter' by @ayulockedin in docs: fix typo 'paramater' -> 'parameter' #5166
• Make nnx.pop remove sown attributes. by @samanklesaria in Make nnx.pop remove sown attributes. #5133
• Rename sharding_names to sharding_metadata by @samanklesaria in Rename sharding_names to sharding_metadata #5089
• Fix bug in graph overhead benchmark by @samanklesaria in Fix bug in graph overhead benchmark #5183
• Fixed typos in the docstrings using antigravity by @vfdev-5 in Fixed typos in the docstrings using antigravity #5145
• docs(nnx): add missing functional args to Conv and LinearGeneral by @ayulockedin in docs(nnx): add missing functional args to Conv and LinearGeneral #5174
• empty change by @copybara-service[bot] in empty change #5184
• Use nnx split during tabulate (clone of Use nnx split during tabulate #5069) by @samanklesaria in Use nnx split during tabulate (clone of #5069) #5186
• Handle pure bodies in nnx.fori_loop by @samanklesaria in Handle pure bodies in nnx.fori_loop #5141
• Typo fix in _cached_partial method by @vfdev-5 in Typo fix in _cached_partial method #5142
• Update mnist example to use NNX (clone of Update mnist example to use NNX #5064) by @samanklesaria in Update mnist example to use NNX (clone of #5064) #5188
• Docs: Fix typo and clarify introduction in Functional API section by @Moriyuki-S in Docs: Fix typo and clarify introduction in Functional API section #5157
• Shard split rngs in lifted vmap if spmd_axis_name is given and applied to vmapped axis. Without this, jax.vmap gives an error: ValueError: Mapped away dimension of inputs passed to vmap should be sharded the same. Got inconsistent axis specs: None vs batch due to split_rngs being replicated. by @copybara-service[bot] in Shard split rngs in lifted vmap if spmd_axis_name is given and applied to vmapped axis. Without this, jax.vmap gives an error: ValueError: Mapped away dimension of inputs passed to vmap should be sharded the same. Got inconsistent axis specs: None vs batch due to split_rngs being replicated. #5189
• Remove ref cycles introduced by self-calling nested functions. by @copybara-service[bot] in Remove ref cycles introduced by self-calling nested functions. #5193
• Add HijaxTransformCoverageTest by @cgarciae in Add HijaxTransformCoverageTest #5190
• allow nnx standalone import p1 by @copybara-service[bot] in allow nnx standalone import p1 #5196
• Add _graph_node_set_key method for List class by @samanklesaria in Add _graph_node_set_key method for List class #5171
• _apply_sharding disallow mixed Explicit/Auto mesh by @copybara-service[bot] in _apply_sharding disallow mixed Explicit/Auto mesh #5199
• update flax to version 0.12.3 by @copybara-service[bot] in update flax to version 0.12.3 #5206

New Contributors

• @ivrolan made their first contribution in fix general guides landing page #5139
• @jackopenn made their first contribution in Add out_sharding arguments to linear layers where supported #5156
• @ayulockedin made their first contribution in docs: fix typo 'paramater' -> 'parameter' #5166
• @Moriyuki-S made their first contribution in Docs: Fix typo and clarify introduction in Functional API section #5157

Full Changelog: v0.12.2...v0.12.3

---

This discussion was created from the release 0.12.3.