Pull Requests google/flax

Add tests for flax.nnx.nn.dtypes

#5275 opened 2026-02-25 22:46 by us

Add PReLU test coverage and move TestPReLUConsistency to activations_test

#5274 opened 2026-02-25 22:07 by us

Fixed two typos in the Module & Pytree guide

#5271 opened 2026-02-23 23:58 by Zirunis

feat(nnx): add out_sharding to recurrent cells

#5255 opened 2026-02-15 14:45 by aarushisingh04

feat(nnx): add preferred_element_type to recurrent cells

#5254 opened 2026-02-15 14:45 by aarushisingh04

feat(nnx): add preferred_element_type to attention module

#5253 opened 2026-02-15 14:45 by aarushisingh04

Track explicit sharding in Variables

#5245 opened 2026-02-12 23:02 by samanklesaria

fix: correct typo 'seperate' to 'separate'

#5229 opened 2026-02-09 04:12 by thecaptain789

enable mlp test

#5224 opened 2026-02-06 03:00 by copybara-service[bot]

Remove nnx.List

#5223 opened 2026-02-04 20:00 by samanklesaria

Add a TAP test for Flax/Keras integration tests.

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

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

#5201 opened 2026-01-26 16:56 by divye-joshi

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 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

[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

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

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

chore: more specific migration guide description

#4977 opened 2025-09-24 09:44 by emergenz