Pull Requests google/flax

Remove nnx.split_rngs call wrapping nnx.scan in linen to nnx tutorial

#5160 opened 2025-12-29 17:15 by samanklesaria

Have _graph_flatten respect nnx.data declarations (extension of #5140)

#5159 opened 2025-12-29 16:42 by samanklesaria

Docs: Fix typo and clarify introduction in Functional API section

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

Add out_sharding arguments to linear layers where supported

#5156 opened 2025-12-25 21:41 by jackopenn

Add out_sharding docstrings to linear layer call methods

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

[pmap] Avoid degraded performance under the new `jax.pmap`.

#5152 opened 2025-12-21 01:48 by copybara-service[bot]

Flax 0.12.2 release

#5150 opened 2025-12-18 20:58 by copybara-service[bot]

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

fix general guides landing page

#5139 opened 2025-12-14 18:39 by ivrolan

Have _graph_flatten respect nnx.data declarations

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

fix kw_only_dataclasses for python 3.14 (part 2)

#5135 opened 2025-12-11 18:34 by copybara-service[bot]

Make `nnx.pop` remove sown attributes.

#5133 opened 2025-12-09 18:30 by samanklesaria

Support multiple None and UNCONSTRAINED when resolving logical rules

#5129 opened 2025-12-08 16:31 by copybara-service[bot]

Add Optimization Cookbook

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

improve hijax guide

#5115 opened 2025-11-25 18:42 by cgarciae

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

Rename sharding_names to sharding_metadata

#5089 opened 2025-11-13 22:12 by samanklesaria

Updated flax_gspmd/bridge_guide to use kernel_metadata, make_mesh and sharding_names

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

Support for python 3.14 pull ready

#5087 opened 2025-11-13 11:20 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