google/flax

Pull Requests Commits

Make Flax thread-safe for multi-threaded tracing.

a-googler committed 2 days ago

db76befd

Make Flax thread-safe for multi-threaded tracing.

BrianWieder committed 2 days ago

c5f2e5b0

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

a-googler committed 3 days ago

abef9d81

dummy-zeros based implementation

mattjj committed 4 days ago

f3b0e126

dce-based solution

mattjj committed 4 days ago

00bd0764

Early exit in Pytree._setattr for data descriptors.

Cristian Garcia committed 7 days ago

3845cd94

add pyrefly suppressions for jax.jit annotation change

Jake VanderPlas committed 8 days ago

c0cb7072

[typing] add assertion to avoid pyrefly error

Jake VanderPlas committed 8 days ago

6d224f1d

Refactor NNX Variable to remove legacy JAX QDD dependencies.

Cristian Garcia committed 14 days ago

c31deb2b

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

a-googler committed 16 days ago

ae2d29e4

implement axes_scan's broadcast semantics via dce

mattjj committed 16 days ago

ba905581

Merge pull request #5481 from samanklesaria:rope

a-googler committed 22 days ago

a5252645

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

a-googler committed 22 days ago

7537c605

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

a-googler committed 24 days ago

92cf9a7d

simplify vmapped metrics reset test

chenkuanliao committed 37 days ago

abc2dbe7

Merge pull request #5526 from samanklesaria:fix_ci_precision

a-googler committed 38 days ago

a31129e2

Use less precise numbers in linen doctest

samanklesaria committed 39 days ago

e2d7ba2d

Automated Code Change

h-joo committed 44 days ago

18f8f311

Add RoPE embeddings

samanklesaria committed 45 days ago

24f467da

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

a-googler committed 46 days ago

b196b016

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

dparikh79 committed 46 days ago

b1b1c184

Support dynamic attributes for Pyrefly in classes using HAS_DYNAMIC_ATTRIBUTES

QEDady committed 48 days ago

07560a7a

merge main into validate-dropout-rate

mohsinm-dev committed 49 days ago

Verified 8f75c1cc

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

a-googler committed 49 days ago

29e55c36

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

a-googler committed 49 days ago

51b3d8f1

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

a-googler committed 49 days ago

067a571d

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

a-googler committed 49 days ago

dd946850

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

chenkuanliao committed 51 days ago

ca14973a

Merge pull request #5511 from google:vfdev-5-patch-1

a-googler committed 53 days ago

dcfabd05

Update tensorflow dependency versions in pyproject.toml

vfdev-5 committed 53 days ago

Verified 395e0331

Older