Flax Files

This release fixes an issue with nnx.List and nnx.Sequential passing all its elements as static when being treated as a pytree.

What's Changed

• fix List and improve Pytree by @cgarciae in https://github.com/google/flax/pull/5072
• Copybara import of the project: by @copybara-service[bot] in https://github.com/google/flax/pull/5211
• Fixed mask sharding if inputs is sharded by @vfdev-5 in https://github.com/google/flax/pull/5212
• nnx.view by @chapman20j in https://github.com/google/flax/pull/5204
• bonsai link by @chapman20j in https://github.com/google/flax/pull/5220
• Add out_sharding argument to Embed layer call by @samanklesaria in https://github.com/google/flax/pull/5205
• Improve Pytree flatten/unflatten by @copybara-service[bot] in https://github.com/google/flax/pull/5216
• Pin optax version to fix CI issues by @samanklesaria in https://github.com/google/flax/pull/5225
• Pin doc deps to lower versions by @samanklesaria in https://github.com/google/flax/pull/5222
• Fix Sequential unwrapping for static attributes by @samanklesaria in https://github.com/google/flax/pull/5218
• feat(nnx): add Grouped Query Attention (GQA) support by @ayulockedin in https://github.com/google/flax/pull/5180
• chore: Migrate gsutil usage to gcloud storage by @gurusai-voleti in https://github.com/google/flax/pull/5230
• Makes linen.Partitioned compatible with linen.WeightNorm by @copybara-service[bot] in https://github.com/google/flax/pull/5234
• maintain data/static definition in split / state by @copybara-service[bot] in https://github.com/google/flax/pull/5227
• Adds a test for linen.WeightNorm and linen.with_partitioning compatibility by @copybara-service[bot] in https://github.com/google/flax/pull/5236
• Add Tunix link by @copybara-service[bot] in https://github.com/google/flax/pull/5239
• Add tree-mode support for nnx.shard_map by @copybara-service[bot] in https://github.com/google/flax/pull/5238
• Rename sharding_names to out_sharding in NNX Variable metadata by @copybara-service[bot] in https://github.com/google/flax/pull/5215

New Contributors

• @gurusai-voleti made their first contribution in https://github.com/google/flax/pull/5230

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