Pull Requests google/flax

Improve Pytree flatten/unflatten

#5216 opened 2026-02-02 23:43 by copybara-service[bot]

Rename sharding_names to out_sharding in NNX Variable metadata

#5215 opened 2026-01-30 23:34 by copybara-service[bot]

Return updates PyTree from nnx.Optimizer.update

#5214 opened 2026-01-28 23:16 by anshuljoshi

Add a TAP test for Flax/Keras integration tests.

#5210 opened 2026-01-28 21:51 by copybara-service[bot]

Add out_sharding argument to Embed layer call

#5205 opened 2026-01-27 17:50 by samanklesaria

feat(nnx): add object-oriented pooling layers and GlobalAveragePool

#5201 opened 2026-01-26 16:56 by starryendymion

optax docs test

#5197 opened 2026-01-21 20:32 by cgarciae

allow nnx standalone import pull ready

#5195 opened 2026-01-20 22:30 by cgarciae

This CL relaxes the input type hints to resolve an invariance issue with Python type checkers (Pyright/MyPy).

#5187 opened 2026-01-13 18:30 by copybara-service[bot]

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

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

feat(nnx): add preferred_element_type to attention.py and recurrent.py, out_sharding to recurrent.py

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

set_mode guide

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

docs(nnx): standardize 'rngs' docstring description

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

[graphcast] Prepare for `jax_pmap_shmap_merge=True`.

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

Add out_sharding docstrings to linear layer call methods

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

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

Migrate VAE example to Flax NNX

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

add support for hijax Variables in nnx transforms

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

Add split and key methods to RngStream and Rngs pull ready

#5047 opened 2025-10-23 14:40 by samanklesaria

Box support

#5029 opened 2025-10-14 23:29 by cgarciae

New set mode

#5028 opened 2025-10-14 21:12 by chapman20j

Make Linen and NNX logical rule deduction align

#5003 opened 2025-10-07 01:20 by IvyZX