google/flax

Pull Requests Commits

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

yashk2810 committed 75 days ago

2d791bb1

Migrate all usages of HiPrimitive to VJPHiPrimitive and delete HiPrimitive

yashk2810 committed 75 days ago

d2b105f0

verify Flax RoPE using mathematical definition instead of PyTorch

IvyZX committed 77 days ago

716b78c4

release Flax 0.12.9

IvyZX committed 77 days ago

2cd6f3c9

Make Flax thread-safe for multi-threaded tracing.

BrianWieder committed 84 days ago

c5f2e5b0

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

a-googler committed 85 days ago

abef9d81

dummy-zeros based implementation

mattjj committed 85 days ago

f3b0e126

dce-based solution

mattjj committed 85 days ago

00bd0764

Early exit in Pytree._setattr for data descriptors.

Cristian Garcia committed 88 days ago

3845cd94

add pyrefly suppressions for jax.jit annotation change

Jake VanderPlas committed 90 days ago

c0cb7072

[typing] add assertion to avoid pyrefly error

Jake VanderPlas committed 90 days ago

6d224f1d

Refactor NNX Variable to remove legacy JAX QDD dependencies.

Cristian Garcia committed 95 days ago

c31deb2b

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

a-googler committed 97 days ago

ae2d29e4

implement axes_scan's broadcast semantics via dce

mattjj committed 97 days ago

ba905581

Merge pull request #5481 from samanklesaria:rope

a-googler committed 104 days ago

a5252645

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

a-googler committed 104 days ago

7537c605

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

a-googler committed 105 days ago

92cf9a7d

simplify vmapped metrics reset test

chenkuanliao committed 119 days ago

abc2dbe7

Merge pull request #5526 from samanklesaria:fix_ci_precision

a-googler committed 120 days ago

a31129e2

Use less precise numbers in linen doctest

samanklesaria committed 120 days ago

e2d7ba2d

Automated Code Change

h-joo committed 125 days ago

18f8f311

Add RoPE embeddings

samanklesaria committed 127 days ago

24f467da

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

a-googler committed 127 days ago

b196b016

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

dparikh79 committed 128 days ago

b1b1c184

Support dynamic attributes for Pyrefly in classes using HAS_DYNAMIC_ATTRIBUTES

QEDady committed 129 days ago

07560a7a

merge main into validate-dropout-rate

mohsinm-dev committed 130 days ago

Verified 8f75c1cc

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

a-googler committed 131 days ago

29e55c36

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

a-googler committed 131 days ago

51b3d8f1

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

a-googler committed 131 days ago

067a571d

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

a-googler committed 131 days ago

dd946850

Older