Commits

Commits on Sep 24, 2026

Release Flax 0.12.10 and use attribute lookup for HiPrim fallback

Cristian Garcia
authored and
Flax Authors
committed
d1f0d69

Commits on Sep 16, 2026

suppress pyrefly errors related to JAX scalar types

Jake VanderPlas
authored and
Flax Authors
committed
4f42ed0

Commits on Aug 21, 2026

[No functional change] Delete VJPHiPrimitive and replace with HiPrim. This is a simple replace.

yashk2810
authored and
Flax Authors
committed
01854da
Migrate all usages of HiPrimitive to VJPHiPrimitive and delete HiPrimitive

yashk2810
authored and
Flax Authors
committed
d2b105f

Commits on Aug 18, 2026

verify Flax RoPE using mathematical definition instead of PyTorch

IvyZX
authored and
Flax Authors
committed
716b78c
release Flax 0.12.9

IvyZX
authored and
Flax Authors
committed
2cd6f3c

Commits on Aug 12, 2026

Make Flax thread-safe for multi-threaded tracing.

BrianWieder
authored and
Flax Authors
committed
c5f2e5b

Commits on Aug 11, 2026

Merge pull request #5551 from mattjj:replace-lazy-init
Flax Authors
committed
abef9d8

Commits on Aug 10, 2026

dummy-zeros based implementation
mattjj
committed
f3b0e12
dce-based solution
mattjj
committed
00bd076

Commits on Aug 7, 2026

Early exit in Pytree._setattr for data descriptors.

Cristian Garcia
authored and
Flax Authors
committed
3845cd9

Commits on Aug 6, 2026

add pyrefly suppressions for jax.jit annotation change

Jake VanderPlas
authored and
Flax Authors
committed
c0cb707
[typing] add assertion to avoid pyrefly error

Jake VanderPlas
authored and
Flax Authors
committed
6d224f1

Commits on Jul 31, 2026

Refactor NNX Variable to remove legacy JAX QDD dependencies.

Cristian Garcia
authored and
Flax Authors
committed
c31deb2

Commits on Jul 29, 2026

Merge pull request #5540 from mattjj:broadcast-via-dce
Flax Authors
committed
ae2d29e
implement axes_scan's broadcast semantics via dce
mattjj
committed
ba90558

Commits on Jul 23, 2026

Merge pull request #5481 from samanklesaria:rope
Flax Authors
committed
a525264
Merge pull request #5503 from mohsinm-dev:validate-dropout-rate
Flax Authors
committed
7537c60

Commits on Jul 21, 2026

Merge pull request #5491 from chenkuanliao:fix-issue-5483
Flax Authors
committed
92cf9a7

Commits on Jul 8, 2026

simplify vmapped metrics reset test
chenkuanliao
committed
abc2dbe

Commits on Jul 7, 2026

Merge pull request #5526 from samanklesaria:fix_ci_precision
Flax Authors
committed
a31129e

Commits on Jul 6, 2026

Use less precise numbers in linen doctest
samanklesaria
committed
e2d7ba2

Commits on Jul 2, 2026

Automated Code Change

h-joo
authored and
Flax Authors
committed
18f8f31

Commits on Jun 30, 2026

Add RoPE embeddings
samanklesaria
committed
24f467d

Commits on Jun 29, 2026

Merge pull request #5475 from dparikh79:fix/5461-nnx-clone-shares-buffers
Flax Authors
committed
b196b01
nnx.clone: add arrays parameter to also copy underlying buffers

dparikh79
authored and
vfdev-5
committed
b1b1c18

Commits on Jun 28, 2026

Support dynamic attributes for Pyrefly in classes using HAS_DYNAMIC_ATTRIBUTES

QEDady
authored and
Flax Authors
committed
07560a7

Commits on Jun 27, 2026

merge main into validate-dropout-rate
mohsinm-dev
committed
8f75c1c

Commits on Jun 26, 2026

Merge pull request #5480 from google:docs-examples--od
Flax Authors
committed
29e55c3
Merge pull request #5505 from mohsinm-dev:fix-attention-assert-to-valueerror
Flax Authors
committed
51b3d8f
Merge pull request #5504 from mohsinm-dev:fix-nnx-promote-dtype-docstring
Flax Authors
committed
067a571
Merge pull request #5479 from google:fix-typos-pth-to-nnx-migration
Flax Authors
committed
dd94685

Commits on Jun 24, 2026

Add vmap reset() regression test for nnx.metrics (#5483)
chenkuanliao
committed
ca14973

Commits on Jun 22, 2026

Merge pull request #5511 from google:vfdev-5-patch-1
Flax Authors
committed
dcfabd0
Update tensorflow dependency versions in pyproject.toml
vfdev-5
authored
395e033