google/flax

Pull Requests Commits

Suppress pyrefly errors related to JAX scalar types

Jake VanderPlas committed 25 days ago

e8f8b45d

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

yashk2810 committed 50 days ago

01854da1

Migrate all usages of HiPrimitive to VJPHiPrimitive and delete HiPrimitive

yashk2810 committed 50 days ago

d2b105f0

verify Flax RoPE using mathematical definition instead of PyTorch

IvyZX committed 52 days ago

716b78c4

release Flax 0.12.9

IvyZX committed 53 days ago

2cd6f3c9

Make Flax thread-safe for multi-threaded tracing.

BrianWieder committed 59 days ago

c5f2e5b0

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

a-googler committed 60 days ago

abef9d81

dummy-zeros based implementation

mattjj committed 60 days ago

f3b0e126

dce-based solution

mattjj committed 60 days ago

00bd0764

Early exit in Pytree._setattr for data descriptors.

Cristian Garcia committed 63 days ago

3845cd94

add pyrefly suppressions for jax.jit annotation change

Jake VanderPlas committed 65 days ago

c0cb7072

[typing] add assertion to avoid pyrefly error

Jake VanderPlas committed 65 days ago

6d224f1d

Refactor NNX Variable to remove legacy JAX QDD dependencies.

Cristian Garcia committed 70 days ago

c31deb2b

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

a-googler committed 72 days ago

ae2d29e4

implement axes_scan's broadcast semantics via dce

mattjj committed 73 days ago

ba905581

Merge pull request #5481 from samanklesaria:rope

a-googler committed 79 days ago

a5252645

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

a-googler committed 79 days ago

7537c605

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

a-googler committed 80 days ago

92cf9a7d

simplify vmapped metrics reset test

chenkuanliao committed 94 days ago

abc2dbe7

Merge pull request #5526 from samanklesaria:fix_ci_precision

a-googler committed 95 days ago

a31129e2

Use less precise numbers in linen doctest

samanklesaria committed 95 days ago

e2d7ba2d

Automated Code Change

h-joo committed 100 days ago

18f8f311

Add RoPE embeddings

samanklesaria committed 102 days ago

24f467da

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

a-googler committed 102 days ago

b196b016

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

dparikh79 committed 103 days ago

b1b1c184

Support dynamic attributes for Pyrefly in classes using HAS_DYNAMIC_ATTRIBUTES

QEDady committed 104 days ago

07560a7a

merge main into validate-dropout-rate

mohsinm-dev committed 105 days ago

Verified 8f75c1cc

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

a-googler committed 106 days ago

29e55c36

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

a-googler committed 106 days ago

51b3d8f1

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

a-googler committed 106 days ago

067a571d

Older