google/flax

Pull Requests Commits

[flax] Widen `LogicalNames` type to accept multi-axis tuples

adityagupta1089 committed 50 minutes ago

71c6ca87

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

yashk2810 committed 10 days ago

01854da1

Migrate all usages of HiPrimitive to VJPHiPrimitive and delete HiPrimitive

yashk2810 committed 10 days ago

d2b105f0

verify Flax RoPE using mathematical definition instead of PyTorch

IvyZX committed 12 days ago

716b78c4

release Flax 0.12.9

IvyZX committed 12 days ago

2cd6f3c9

Make Flax thread-safe for multi-threaded tracing.

BrianWieder committed 19 days ago

c5f2e5b0

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

a-googler committed 20 days ago

abef9d81

dummy-zeros based implementation

mattjj committed 20 days ago

f3b0e126

dce-based solution

mattjj committed 20 days ago

00bd0764

Early exit in Pytree._setattr for data descriptors.

Cristian Garcia committed 23 days ago

3845cd94

add pyrefly suppressions for jax.jit annotation change

Jake VanderPlas committed 25 days ago

c0cb7072

[typing] add assertion to avoid pyrefly error

Jake VanderPlas committed 25 days ago

6d224f1d

Refactor NNX Variable to remove legacy JAX QDD dependencies.

Cristian Garcia committed 30 days ago

c31deb2b

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

a-googler committed 32 days ago

ae2d29e4

implement axes_scan's broadcast semantics via dce

mattjj committed 32 days ago

ba905581

Merge pull request #5481 from samanklesaria:rope

a-googler committed 39 days ago

a5252645

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

a-googler committed 39 days ago

7537c605

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

a-googler committed 40 days ago

92cf9a7d

simplify vmapped metrics reset test

chenkuanliao committed 54 days ago

abc2dbe7

Merge pull request #5526 from samanklesaria:fix_ci_precision

a-googler committed 55 days ago

a31129e2

Use less precise numbers in linen doctest

samanklesaria committed 55 days ago

e2d7ba2d

Automated Code Change

h-joo committed 60 days ago

18f8f311

Add RoPE embeddings

samanklesaria committed 61 days ago

24f467da

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

a-googler committed 62 days ago

b196b016

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

dparikh79 committed 62 days ago

b1b1c184

Support dynamic attributes for Pyrefly in classes using HAS_DYNAMIC_ATTRIBUTES

QEDady committed 64 days ago

07560a7a

merge main into validate-dropout-rate

mohsinm-dev committed 65 days ago

Verified 8f75c1cc

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

a-googler committed 66 days ago

29e55c36

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

a-googler committed 66 days ago

51b3d8f1

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

a-googler committed 66 days ago

067a571d

Older