Pull Requests google/flax

Reduce warning message of about GCS usage

#5181 opened 2026-01-08 17:21 by wdhongtw

feat(nnx): add Grouped Query Attention (GQA) support

#5180 opened 2026-01-07 23:23 by ayulockedin

feat(nnx): add preferred_element_type and out_sharding support to Attention and RNN modules

#5179 opened 2026-01-07 21:09 by aarushisingh04

set_mode guide

#5178 opened 2026-01-07 20:23 by chapman20j

docs(nnx): add missing functional args to Conv and LinearGeneral

#5174 opened 2026-01-06 18:51 by ayulockedin

docs(nnx): standardize 'rngs' docstring description

#5173 opened 2026-01-06 18:29 by ayulockedin

Add _graph_node_set_key method for List class

#5171 opened 2026-01-06 16:34 by samanklesaria

fix: remove superfluous `jax.Ref` standalone statement

#5169 opened 2026-01-06 15:16 by pfackeldey

[graphcast] Prepare for `jax_pmap_shmap_merge=True`.

#5168 opened 2026-01-06 06:46 by copybara-service[bot]

refactor: add missing type hints to flax/core/scope.py

#5167 opened 2026-01-05 23:55 by ayulockedin

Docs: Fix typo and clarify introduction in Functional API section

#5157 opened 2025-12-26 10:23 by Moriyuki-S

Add out_sharding docstrings to linear layer call methods

#5154 opened 2025-12-25 21:04 by jackopenn

Fixed typos in the docstrings using antigravity

#5145 opened 2025-12-18 14:05 by vfdev-5

Typo fix in _cached_partial method

#5142 opened 2025-12-18 08:55 by vfdev-5

Handle pure bodies in nnx.fori_loop

#5141 opened 2025-12-17 19:29 by samanklesaria

flatten respect nnx.data

#5140 opened 2025-12-15 17:13 by cgarciae

Have _graph_flatten respect nnx.data declarations

#5136 opened 2025-12-12 20:19 by samanklesaria

Add Optimization Cookbook

#5117 opened 2025-11-28 20:02 by samanklesaria

Added sharding propagation support in nnx.eval_shape

#5111 opened 2025-11-24 15:37 by vfdev-5

Rewritten flax.jax_utils.prefetch_to_device and flax.jax_utils.replicate using jax.device_put pull ready

#5101 opened 2025-11-19 17:47 by vfdev-5

Reenabled examples tests

#5098 opened 2025-11-18 10:58 by vfdev-5

Refactor nnx.tabulate

#5096 opened 2025-11-18 06:07 by cgarciae

Fixed dataset path in train tests in examples

#5092 opened 2025-11-17 10:19 by vfdev-5

Updated flax_gspmd/bridge_guide to use kernel_metadata, make_mesh and sharding_names

#5088 opened 2025-11-13 15:52 by vfdev-5

Updated pyproject.toml

#5086 opened 2025-11-13 10:45 by vfdev-5

Migrate VAE example to Flax NNX

#5077 opened 2025-11-07 20:54 by sanepunk

fix List and improve Pytree

#5072 opened 2025-11-04 03:10 by cgarciae

Use nnx split during `tabulate` pull ready

#5069 opened 2025-11-03 21:50 by samanklesaria

add support for hijax Variables in nnx transforms

#5065 opened 2025-11-01 19:19 by cgarciae

Update mnist example to use NNX pull ready

#5064 opened 2025-10-31 17:53 by samanklesaria