Pull Requests google/flax

fix cached_partial for raw array attributes

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

fix MultiMetric GraphDef not hashable under jit

#5492 opened 2026-06-11 02:22 by mohsinm-dev

Make metric reset() shape-preserving under vmap (#5483)

#5491 opened 2026-06-11 00:26 by chenkuanliao

Handle capture variables in `nnx.jit`.

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

WIP: 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]

feat(nnx): add arrays=True to nnx.clone for independent buffer copies

#5482 opened 2026-06-04 16:51 by Sumu004

Add RoPE embeddings

#5481 opened 2026-06-02 17:25 by samanklesaria

Added Object Detection example with minimal DETR pull ready

#5480 opened 2026-06-02 13:51 by vfdev-5

Fixed typos in docs_nnx/migrating/pytorch_to_jax_flax.rst

#5479 opened 2026-06-02 12:24 by vfdev-5

optimize jit_partial

#5476 opened 2026-05-31 03:04 by copybara-service[bot]

nnx.clone: add arrays parameter to also copy underlying buffers

#5475 opened 2026-05-29 21:19 by dparikh79

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]

feat(nnx): add DyT (Dynamic Tanh) normalization module

#5459 opened 2026-05-11 13:17 by Ashutosh0x

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

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

Added ViT training example pull ready

#5455 opened 2026-05-05 14:52 by vfdev-5

Bump the pip group across 10 directories with 4 updates dependencies python

#5444 opened 2026-04-24 18:17 by dependabot[bot]

migrate to pyrefly

#5437 opened 2026-04-22 14:16 by copybara-service[bot]

Bump the uv group across 10 directories with 5 updates dependencies python:uv

#5432 opened 2026-04-18 00:20 by dependabot[bot]

Avoid usage of deprecated `jax.core` APIs.

#5420 opened 2026-04-15 18:52 by copybara-service[bot]

use State instead of GraphState

#5401 opened 2026-04-08 17:26 by copybara-service[bot]

use_running_average and deterministic default to None

#5399 opened 2026-04-07 23:35 by copybara-service[bot]

Consolidate "view" functions with similar naming convention.

#5376 opened 2026-03-25 21:09 by samanklesaria

Upgrade third_party/python_runtime/v3_13_unstable to 3.13.11

#5371 opened 2026-03-24 09:54 by copybara-service[bot]

Remove spmd_lowering from MeshComputation and do more cleanups after pmap, PmapSharding is gone

#5360 opened 2026-03-20 06:09 by copybara-service[bot]

improve scaling up guide

#5354 opened 2026-03-18 19:30 by cgarciae

Rename `SDS(vma: frozenset = ...)` to `SDS(manual_type: jax.sharding.ManualAxisType = ...)`

#5351 opened 2026-03-18 02:16 by copybara-service[bot]