Pull Requests google/flax

Add NNX WeightNorm (update of #4568) pull ready

#5043 by samanklesaria was merged 2025-10-24 14:44

Recursive map pull ready

#5042 by chapman20j was merged 2025-10-24 19:29

Copybara import of the project:

#5041 by copybara-service[bot] was merged 2025-10-21 18:24

Recursive Map pull ready

#5040 by chapman20j was closed 2025-10-21 21:51

Fix shard_map documentation link in compilation.py pull ready

#5038 by vfdev-5 was merged 2025-10-24 15:00

[docs] Exposed more helper functions/classes in state.rst pull ready

#5037 by vfdev-5 was merged 2025-10-21 18:22

Add REINFORCE policy gradient example for CartPole using Flax NNX

#5036 by sanepunk was closed 2025-10-31 13:10

Deduplicate `Variable` nodes in `iter_graph` and eliminate recursion.

#5035 by copybara-service[bot] was merged 2025-10-20 16:10

fix reseed for abstract values pull ready

#5034 by cgarciae was merged 2025-10-17 20:42

Add flip doc link to eager sharding error message pull ready

#5033 by IvyZX was merged 2025-10-16 23:20

Support for python 3.14 pull ready

#5032 by vfdev-5 was merged 2025-10-21 16:22

Added a note in nnx.jit about arg donation pull ready

#5031 by vfdev-5 was merged 2025-10-16 23:15

Fix formatting in PR template checklist pull ready

#5024 by rapsealk was merged 2025-10-15 21:23

Update README to use fully qualified `nnx.Linear` in example pull ready

#5023 by rapsealk was merged 2025-10-15 21:25

Fixed attribute visualization in __treescope_repr__ pull ready

#5022 by vfdev-5 was merged 2025-10-15 21:24

Use jax.shard_map in flax

#5020 by copybara-service[bot] was merged 2025-10-11 02:35

deprecate .value property

#5019 by cgarciae was closed 2025-10-13 20:20

use jax.shard_map

#5018 by copybara-service[bot] was merged 2025-10-11 02:54

Add Rngs KeylessInitializers pull ready

#5017 by cgarciae was merged 2025-10-10 00:37

Remove invalid gymnasium dependency in pyproject.toml pull ready

#5016 by IvyZX was merged 2025-10-10 01:16

optimize scan transpositions pull ready

#5015 by cgarciae was merged 2025-10-10 00:40

[flax:examples:vae] Small linter fixes.

#5014 by copybara-service[bot] was merged 2025-12-01 21:40

[flax:examples:gemma] Fixing linter errors.

#5013 by copybara-service[bot] was merged 2025-12-01 22:07

[flax:examples:wmt] Small linter fixes.

#5012 by copybara-service[bot] was merged 2025-12-01 20:12

Migrate from `jax.experimental.enable_x64` to `jax.enable_x64`.

#5011 by copybara-service[bot] was merged 2025-10-09 23:12

Use sphinx 6.2+ for docs, which works with Python 3.13. pull ready

#5009 by jburnim was merged 2025-10-09 21:09

Fix nnx tabulate variable hooks pull ready

#5008 by mohsinm-dev was merged 2025-10-15 21:28

feat: add `nnx.set_metadata` to in-place change metadata of the state variables of `nnx.Module`s pull ready

#5007 by pfackeldey was merged 2025-10-15 21:25

Variable refactor pull ready

#5006 by cgarciae was merged 2025-10-10 01:14

[flax:examples:mnist] Update mnist example to use NNX.

#5002 by copybara-service[bot] was closed 2025-10-31 18:14