Pull Requests google/flax

Ported nnx.PReLU from linen pull ready

#4934 by vfdev-5 was merged 2025-09-12 19:09

fix mypy

#4931 by copybara-service[bot] was merged 2025-09-09 14:48

Added nnx.scan docs and few minor docs fixes pull ready

#4930 by vfdev-5 was merged 2025-09-15 17:58

finish pytree guide pull ready

#4929 by cgarciae was merged 2025-09-18 23:41

Set mode

#4927 by cgarciae was closed 2025-11-10 08:53

Avoid passing non-boolean mask to `where` argument of `jax.numpy` reductions. Non-boolean mask inputs have been deprecated for several releases, and will result in an error starting in JAX v0.8.0.

#4923 by copybara-service[bot] was merged 2025-09-09 15:47

Added missing functions to graph.rst pull ready

#4922 by vfdev-5 was merged 2025-09-02 21:33

Added preferred_element_type arg to nnx.Linear*, nnx.Conv*, nnx.Einsum pull ready

#4920 by vfdev-5 was merged 2025-09-03 21:55

Update flax/docs_nnx/guides/performance.md and .ipynb pull ready

#4919 by hanrach9 was merged 2025-09-03 16:52

Fix typo in unflatten docs

#4918 by copybara-service[bot] was merged 2025-08-28 21:41

fix RNN

#4917 by copybara-service[bot] was merged 2025-08-28 23:12

update version to 0.11.2

#4915 by copybara-service[bot] was merged 2025-08-28 17:25

Doc fix: remove dead link to pre-Orbax checkpointing.

#4914 by copybara-service[bot] was merged 2025-08-28 21:23

merge no copy Variables pull ready

#4912 by cgarciae was merged 2025-08-27 18:43

[jax:benchmarks] Add tracing/lowering benchmarks for a few flax examples.

#4911 by copybara-service[bot] was merged 2025-12-09 00:26

Doc fix: remove outdated advice on flax v0.6.10; it was released two years ago.

#4910 by copybara-service[bot] was merged 2025-08-26 19:41

fix mypy on main pull ready

#4909 by cgarciae was merged 2025-08-26 21:32

Update README badges and remove invalid ones pull ready

#4905 by IvyZX was merged 2025-09-03 22:35

Update optimizer.py to support masked variable from optax. pull ready

#4904 by ywrt was merged 2025-09-02 15:11

revert is_leaf logic in _check_carry_same_references

#4903 by copybara-service[bot] was merged 2025-08-26 00:03

document jax.random shorthand methods pull ready

#4899 by cgarciae was merged 2025-08-20 12:37

Fix bug when raising ScopeParamNotFoundError.

#4898 by copybara-service[bot] was merged 2025-08-26 20:43

static + pytree guide pull ready

#4897 by cgarciae was merged 2025-09-05 00:15

add find_duplicates pull ready

#4894 by cgarciae was merged 2025-08-18 22:04

Sharding API improvements (non breaking) pull ready

#4893 by IvyZX was merged 2025-08-19 22:33

Pin sentencepiece version to 0.2.0 to fix head pull ready

#4892 by IvyZX was merged 2025-08-14 23:07

Fix nnx.tabulate crash with empty dict/None values (fixes #4889) pull ready

#4891 by mohsinm-dev was merged 2025-09-29 18:29

call jax's source_info_util.register_exclusion in flax's traceback_util.register_exclusion

#4887 by copybara-service[bot] was merged 2025-08-13 22:27

Internal change

#4886 by copybara-service[bot] was merged 2025-08-12 22:08

Update doc build instructions and clean up unused packages pull ready

#4885 by IvyZX was merged 2025-08-12 00:35