Pull Requests google/flax

Refactor NNX Variable to remove legacy JAX QDD dependencies.

#5541 opened 2026-07-30 21:29 by copybara-service[bot]

Fix vmap axes for input_positions in dot_product_attention_with_rope

#5537 opened 2026-07-26 20:26 by TheSaiEaranti

Fix broken links

#5533 opened 2026-07-22 20:32 by copybara-service[bot]

Exclude haiku doctests from CI

#5532 opened 2026-07-22 19:37 by samanklesaria

Handle captures in jax transforms with hijax

#5531 opened 2026-07-21 20:28 by samanklesaria

Preserve RngStream counters for passthrough streams in Rngs.split/broadcast

#5529 opened 2026-07-17 17:22 by sohumt123

Fix nnx.Dropout crash when broadcast_dims includes a sharded axis

#5528 opened 2026-07-17 17:21 by sohumt123

Handle captures in nnx.vmap

#5527 opened 2026-07-07 00:40 by samanklesaria

Fix tracing cache misses in nnx.cond and nnx.switch (#5512)

#5525 opened 2026-07-06 14:25 by kapellirohith

Add pickle/cloudpickle round-trip tests for Linen modules (#1481)

#5524 opened 2026-07-06 12:44 by kapellirohith

Speed up NNX graph flattening by avoiding unnecessary array-ref checks

#5523 opened 2026-07-06 05:24 by KookiesNKareem

Handle captures grad

#5522 opened 2026-07-02 21:15 by samanklesaria

Update NNX docs for deprecated Variable.value access

#5521 opened 2026-07-02 07:23 by sumitdotml

docs: clarify shape dynamics and add example for make_causal_mask

#5519 opened 2026-06-27 23:06 by SakshieP

Fix nnx.cond tracing cache misses by caching SimpleCondFn wrappers

#5518 opened 2026-06-27 16:46 by chenkuanliao

cache cond/switch branch wrappers to avoid tracing cache misses

#5517 opened 2026-06-27 08:38 by mohsinm-dev

Handle captures in nnx.scan

#5515 opened 2026-06-25 18:52 by samanklesaria

Optimize Flax ogbg_molpcba example GNN model with custom Pallas Fused LayerNorm kernel, custom VJP for reverse-mode autodiff, and numerical correctness benchmark script.

#5513 opened 2026-06-22 19:40 by copybara-service[bot]

fix(attention): validate dropout_rate range and require dropout_rng w…

#5510 opened 2026-06-21 04:39 by Vitaliy-Pikalo

default partition metadata in vmap/scan/pmap

#5508 opened 2026-06-18 02:52 by mohsinm-dev

add nnx:v1

#5502 opened 2026-06-16 00:08 by copybara-service[bot]

fix cached_partial for raw array attributes

#5501 opened 2026-06-13 11:15 by discobot

Handle capture variables in `nnx.jit`.

#5490 opened 2026-06-09 21:43 by samanklesaria

Add multimodal llm example

#5488 opened 2026-06-09 20:28 by samanklesaria

allow captures in nnx.grad

#5484 opened 2026-06-04 20:00 by copybara-service[bot]

Add docstrings to nnx filterlib filter classes

#5474 opened 2026-05-29 04:16 by santihdzs

Add sepreformer example

#5472 opened 2026-05-26 16:29 by samanklesaria

Fix MGUCell hidden state update

#5467 opened 2026-05-21 05:18 by lanarkite99

Migrate away from private jax._src.tree_util._registry

#5464 opened 2026-05-20 01:35 by copybara-service[bot]

Add pyrefly type checking: pyrefly.toml config, CI integration

#5458 opened 2026-05-11 11:19 by vfdev-5