google/flax

Pull Requests Commits

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

yashk2810 committed 1 day ago

2d791bb1

Migrate all usages of HiPrimitive to VJPHiPrimitive and delete HiPrimitive

yashk2810 committed 1 day ago

d2b105f0

verify Flax RoPE using mathematical definition instead of PyTorch

IvyZX committed 4 days ago

716b78c4

release Flax 0.12.9

IvyZX committed 4 days ago

2cd6f3c9

Make Flax thread-safe for multi-threaded tracing.

BrianWieder committed 10 days ago

c5f2e5b0

Merge pull request #5551 from mattjj:replace-lazy-init

a-googler committed 11 days ago

abef9d81

dummy-zeros based implementation

mattjj committed 12 days ago

f3b0e126

dce-based solution

mattjj committed 12 days ago

00bd0764

Early exit in Pytree._setattr for data descriptors.

Cristian Garcia committed 15 days ago

3845cd94

add pyrefly suppressions for jax.jit annotation change

Jake VanderPlas committed 16 days ago

c0cb7072

[typing] add assertion to avoid pyrefly error

Jake VanderPlas committed 16 days ago

6d224f1d

Refactor NNX Variable to remove legacy JAX QDD dependencies.

Cristian Garcia committed 22 days ago

c31deb2b

Merge pull request #5540 from mattjj:broadcast-via-dce

a-googler committed 24 days ago

ae2d29e4

implement axes_scan's broadcast semantics via dce

mattjj committed 24 days ago

ba905581

Merge pull request #5481 from samanklesaria:rope

a-googler committed 30 days ago

a5252645

Merge pull request #5503 from mohsinm-dev:validate-dropout-rate

a-googler committed 30 days ago

7537c605

Merge pull request #5491 from chenkuanliao:fix-issue-5483

a-googler committed 32 days ago

92cf9a7d

simplify vmapped metrics reset test

chenkuanliao committed 45 days ago

abc2dbe7

Merge pull request #5526 from samanklesaria:fix_ci_precision

a-googler committed 46 days ago

a31129e2

Use less precise numbers in linen doctest

samanklesaria committed 47 days ago

e2d7ba2d

Automated Code Change

h-joo committed 51 days ago

18f8f311

Add RoPE embeddings

samanklesaria committed 53 days ago

24f467da

Merge pull request #5475 from dparikh79:fix/5461-nnx-clone-shares-buffers

a-googler committed 54 days ago

b196b016

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

dparikh79 committed 54 days ago

b1b1c184

Support dynamic attributes for Pyrefly in classes using HAS_DYNAMIC_ATTRIBUTES

QEDady committed 55 days ago

07560a7a

merge main into validate-dropout-rate

mohsinm-dev committed 56 days ago

Verified 8f75c1cc

Merge pull request #5480 from google:docs-examples--od

a-googler committed 57 days ago

29e55c36

Merge pull request #5505 from mohsinm-dev:fix-attention-assert-to-valueerror

a-googler committed 57 days ago

51b3d8f1

Merge pull request #5504 from mohsinm-dev:fix-nnx-promote-dtype-docstring

a-googler committed 57 days ago

067a571d

Merge pull request #5479 from google:fix-typos-pth-to-nnx-migration

a-googler committed 57 days ago

dd946850

Older