Pull Requests google/flax

linen: add TernaryDense — ternary-weight linear layer {-1, 0, +1}

#5466 by eriirfos-eng was closed 2026-05-21 13:32

add storage backend type in saving telemetry.

#5465 by copybara-service[bot] was merged 2026-05-20 17:52

Fixed mypy errors on CI

#5460 by vfdev-5 was closed 2026-05-18 08:26

Add data loaders tutorial pull ready

#5454 by samanklesaria was merged 2026-05-19 23:06

optimize transform updates

#5453 by copybara-service[bot] was merged 2026-05-22 00:05

add nnx.EMA

#5451 by copybara-service[bot] was merged 2026-05-11 18:13

improve Optimizer

#5450 by copybara-service[bot] was merged 2026-04-28 19:05

fix Pytree typing issue

#5449 by copybara-service[bot] was merged 2026-04-28 18:32

Use nnx.prefix rather than StateAxes in split_rngs docstring

#5448 by samanklesaria was closed 2026-04-29 22:00

Add fix for msgpack serialization of nnx.List structures pull ready

#5446 by samanklesaria was merged 2026-04-30 17:01

Fix docstrings in graphlib under graph=False pull ready

#5443 by samanklesaria was merged 2026-04-29 22:10

Fix rnglib doctests for graph_updates=False default pull ready

#5442 by samanklesaria was merged 2026-05-12 22:41

Fix doctests in transforms for graph_updates=False default pull ready

#5441 by samanklesaria was merged 2026-05-12 22:48

Fix update doctest for graph_updates=False default pull ready

#5440 by samanklesaria was merged 2026-05-12 22:43

Fix doctests for cached_partial in tree mode

#5439 by samanklesaria was closed 2026-04-30 16:34

Disallow `P.UNCONSTRAINED` in `out_shardings` in `jit`. This leads to simplifications in how we lower shardings to stableHLO.

#5435 by copybara-service[bot] was merged 2026-04-22 02:15

Added out_sharding and qkv_sharding to MHA call method pull ready

#5434 by vfdev-5 was merged 2026-05-11 16:16

Add machine translation example pull ready

#5431 by samanklesaria was merged 2026-04-23 21:02

Avoid usage of deprecated `jax.core` APIs.

#5430 by copybara-service[bot] was merged 2026-04-18 00:16

update version to 0.12.7

#5429 by copybara-service[bot] was merged 2026-04-19 16:10

Avoid usage of deprecated `jax.core` APIs.

#5428 by copybara-service[bot] was closed 2026-04-21 18:06

Updated gemma docs pull ready

#5427 by vfdev-5 was merged 2026-04-23 21:08

Avoid usage of deprecated `jax.core` APIs.

#5426 by copybara-service[bot] was closed 2026-04-18 00:17

Update mnist tutorial to use tensorboard pull ready

#5425 by samanklesaria was merged 2026-05-12 22:36

Reorganize examples pull ready

#5424 by samanklesaria was merged 2026-04-23 20:30

Move diffusion example to other examples directory

#5423 by samanklesaria was closed 2026-04-16 17:12

Fix doubled 'the the' typos across multiple files pull ready

#5422 by SAY-5 was merged 2026-04-21 23:11

Added kwargs to call to pass kwargs to basemodule in nnx.LoRA pull ready

#5421 by vfdev-5 was merged 2026-04-17 19:12

Avoid usage of deprecated `jax.core` APIs.

#5419 by copybara-service[bot] was closed 2026-04-19 16:10

add nnx.unpack

#5418 by copybara-service[bot] was merged 2026-04-15 00:07