Pull Requests google/flax

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]

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

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

chore: more specific migration guide description

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

Add a customizable conv_transpose. pull ready

#4938 opened 2025-09-11 21:52 by copybara-service[bot]

Don't pass linen_meta_type argument when creating AxisMetadata subclasses.

#4933 opened 2025-09-09 11:42 by copybara-service[bot]

Improve error message with class name.

#4932 opened 2025-09-09 08:38 by copybara-service[bot]

adds recursive_map

#4926 opened 2025-09-05 01:08 by copybara-service[bot]