Pull Requests google/flax

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

Avoid install, import, or tests with tensorflow-text under Python 3.13+. pull ready

#5001 by jburnim was merged 2025-10-07 19:03

Convert linen pytorch guide to nnx pull ready

#4999 by samanklesaria was merged 2025-10-27 20:03

Removed kernel_init/bias_init atttributes from popular layers pull ready

#4998 by vfdev-5 was merged 2025-10-09 22:22

Fix NNX jit static args with in_shardings issue #4989 pull ready

#4996 by mohsinm-dev was merged 2025-11-06 17:09

disallow setting metadata through settattr pull ready

#4993 by cgarciae was merged 2025-10-07 22:14

Fixed broken nnx.statelib.diff pull ready

#4992 by vfdev-5 was merged 2025-10-01 01:48

Use `jnp.stack` instead of `np.stack` in `flax.training.common_utils.stack_forest` pull ready

#4991 by vfdev-5 was merged 2025-10-01 01:23

Expand ConvTranspose padding documentation pull ready

#4990 by samanklesaria was merged 2025-11-12 01:52

clean up jax.Ref-related names

#4988 by copybara-service[bot] was merged 2025-09-29 17:48

python 3.13 support pull ready

#4987 by cgarciae was merged 2025-10-16 19:22

Future-proof imports of jax.new_ref / jax.Ref.

#4986 by copybara-service[bot] was merged 2025-09-29 20:08

Improve Variable.{get,set}_metadata pull ready

#4985 by cgarciae was merged 2025-10-01 23:35

propagate eager sharding pull ready

#4983 by cgarciae was merged 2025-09-25 23:43

update to version 0.12.0 pull ready

#4982 by cgarciae was merged 2025-09-25 19:37

[nnx] preserve the function's type information in jit pull ready

#4981 by cgarciae was merged 2025-09-25 22:37

Move `ArrayRef` creation to the end of `Variable` creation pull ready

#4980 by IvyZX was merged 2025-09-26 13:27

Copybara import of the project:

#4976 by copybara-service[bot] was merged 2025-09-24 20:19

Update tutorial examples to thread explicit RNGs pull ready

#4975 by samanklesaria was merged 2025-10-30 22:40

ignore uv.lock

#4974 by copybara-service[bot] was merged 2025-09-25 22:02

Minor typo fixes in flax gspmd guide pull ready

#4970 by vfdev-5 was merged 2025-09-25 20:04

Correctly expose `flax.config.temp_flip_flag` pull ready

#4969 by IvyZX was merged 2025-09-19 23:55

add Variable.set_metadata pull ready

#4968 by cgarciae was merged 2025-09-25 22:39

raise on Variable inplace operators pull ready

#4967 by cgarciae was merged 2025-09-24 20:06

import PathParts from flax.typing pull ready

#4966 by cgarciae was merged 2025-09-19 23:53