Pull Requests google/flax

Fixed mypy errors on CI

#5460 opened 2026-05-11 13:46 by vfdev-5

feat(nnx): add DyT (Dynamic Tanh) normalization module

#5459 opened 2026-05-11 13:17 by Ashutosh0x

Add pyrefly type checking: pyrefly.toml config, CI integration

#5458 opened 2026-05-11 11:19 by vfdev-5

Added ViT training example

#5455 opened 2026-05-05 14:52 by vfdev-5

Add data loaders tutorial

#5454 opened 2026-05-01 21:36 by samanklesaria

optimize transform updates

#5453 opened 2026-04-30 22:16 by copybara-service[bot]

Bump the pip group across 10 directories with 4 updates dependencies python

#5444 opened 2026-04-24 18:17 by dependabot[bot]

Fix rnglib doctests for graph_updates=False default

#5442 opened 2026-04-24 15:44 by samanklesaria

Fix doctests in transforms for graph_updates=False default

#5441 opened 2026-04-24 15:43 by samanklesaria

Fix update doctest for graph_updates=False default

#5440 opened 2026-04-24 15:42 by samanklesaria

migrate to pyrefly

#5437 opened 2026-04-22 14:16 by copybara-service[bot]

Bump the uv group across 10 directories with 5 updates dependencies python:uv

#5432 opened 2026-04-18 00:20 by dependabot[bot]

Update mnist tutorial to use tensorboard

#5425 opened 2026-04-17 02:41 by samanklesaria

Avoid usage of deprecated `jax.core` APIs.

#5420 opened 2026-04-15 18:52 by copybara-service[bot]

use State instead of GraphState

#5401 opened 2026-04-08 17:26 by copybara-service[bot]

use_running_average and deterministic default to None

#5399 opened 2026-04-07 23:35 by copybara-service[bot]

Consolidate "view" functions with similar naming convention.

#5376 opened 2026-03-25 21:09 by samanklesaria

Upgrade third_party/python_runtime/v3_13_unstable to 3.13.11

#5371 opened 2026-03-24 09:54 by copybara-service[bot]

Update the RNG guide with split and without StateAxes

#5363 opened 2026-03-20 17:06 by samanklesaria

Remove spmd_lowering from MeshComputation and do more cleanups after pmap, PmapSharding is gone

#5360 opened 2026-03-20 06:09 by copybara-service[bot]

improve scaling up guide

#5354 opened 2026-03-18 19:30 by cgarciae

Rename `SDS(vma: frozenset = ...)` to `SDS(manual_type: jax.sharding.ManualAxisType = ...)`

#5351 opened 2026-03-18 02:16 by copybara-service[bot]

Add 'prefix' argument to temp_flip_flag function

#5343 opened 2026-03-17 20:28 by samanklesaria

Fix broken error-message formatting in flax.nnx

#5334 opened 2026-03-16 13:35 by gdrosa

add transforms_tree guide

#5307 opened 2026-03-04 18:22 by cgarciae

add nnx_basics_tree.ipynb

#5293 opened 2026-03-02 19:16 by cgarciae

Expand Einsum test coverage with parameterized tests

#5285 opened 2026-02-26 21:35 by us

Add tests, docstrings and validation to metrics.py

#5284 opened 2026-02-26 21:33 by us

Add examples and fix docstrings in nnx.nn.dtypes

#5283 opened 2026-02-26 21:11 by us

Expand Dropout test coverage and migrate to absltest

#5281 opened 2026-02-26 20:32 by us