Flax Files

What's Changed

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

New Contributors

• @ivrolan made their first contribution in https://github.com/google/flax/pull/5139
• @jackopenn made their first contribution in https://github.com/google/flax/pull/5156
• @ayulockedin made their first contribution in https://github.com/google/flax/pull/5166
• @Moriyuki-S made their first contribution in https://github.com/google/flax/pull/5157

Full Changelog: https://github.com/google/flax/compare/v0.12.2...v0.12.3