google/flax

Pull Requests Commits

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

adityagupta1089 committed 23 days ago

bd3aacd4

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

yashk2810 committed 33 days ago

01854da1

Migrate all usages of HiPrimitive to VJPHiPrimitive and delete HiPrimitive

yashk2810 committed 34 days ago

d2b105f0

verify Flax RoPE using mathematical definition instead of PyTorch

IvyZX committed 36 days ago

716b78c4

release Flax 0.12.9

IvyZX committed 36 days ago

2cd6f3c9

Make Flax thread-safe for multi-threaded tracing.

BrianWieder committed 43 days ago

c5f2e5b0

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

a-googler committed 43 days ago

abef9d81

dummy-zeros based implementation

mattjj committed 44 days ago

f3b0e126

dce-based solution

mattjj committed 44 days ago

00bd0764

Early exit in Pytree._setattr for data descriptors.

Cristian Garcia committed 47 days ago

3845cd94

add pyrefly suppressions for jax.jit annotation change

Jake VanderPlas committed 49 days ago

c0cb7072

[typing] add assertion to avoid pyrefly error

Jake VanderPlas committed 49 days ago

6d224f1d

Refactor NNX Variable to remove legacy JAX QDD dependencies.

Cristian Garcia committed 54 days ago

c31deb2b

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

a-googler committed 56 days ago

ae2d29e4

implement axes_scan's broadcast semantics via dce

mattjj committed 56 days ago

ba905581

Merge pull request #5481 from samanklesaria:rope

a-googler committed 62 days ago

a5252645

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

a-googler committed 63 days ago

7537c605

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

a-googler committed 64 days ago

92cf9a7d

simplify vmapped metrics reset test

chenkuanliao committed 77 days ago

abc2dbe7

Merge pull request #5526 from samanklesaria:fix_ci_precision

a-googler committed 78 days ago

a31129e2

Use less precise numbers in linen doctest

samanklesaria committed 79 days ago

e2d7ba2d

Automated Code Change

h-joo committed 84 days ago

18f8f311

Add RoPE embeddings

samanklesaria committed 85 days ago

24f467da

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

a-googler committed 86 days ago

b196b016

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

dparikh79 committed 86 days ago

b1b1c184

Support dynamic attributes for Pyrefly in classes using HAS_DYNAMIC_ATTRIBUTES

QEDady committed 88 days ago

07560a7a

merge main into validate-dropout-rate

mohsinm-dev committed 89 days ago

Verified 8f75c1cc

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

a-googler committed 89 days ago

29e55c36

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

a-googler committed 89 days ago

51b3d8f1

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

a-googler committed 89 days ago

067a571d

Older