google/flax

Pull Requests Commits

Early exit in Pytree._setattr for data descriptors.

Cristian Garcia committed 26 minutes ago

3845cd94

add pyrefly suppressions for jax.jit annotation change

Jake VanderPlas committed 1 day ago

c0cb7072

[typing] add assertion to avoid pyrefly error

Jake VanderPlas committed 1 day ago

6d224f1d

Refactor NNX Variable to remove legacy JAX QDD dependencies.

Cristian Garcia committed 7 days ago

c31deb2b

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

a-googler committed 9 days ago

ae2d29e4

implement axes_scan's broadcast semantics via dce

mattjj committed 9 days ago

ba905581

Merge pull request #5481 from samanklesaria:rope

a-googler committed 15 days ago

a5252645

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

a-googler committed 15 days ago

7537c605

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

a-googler committed 17 days ago

92cf9a7d

simplify vmapped metrics reset test

chenkuanliao committed 30 days ago

abc2dbe7

Merge pull request #5526 from samanklesaria:fix_ci_precision

a-googler committed 31 days ago

a31129e2

Use less precise numbers in linen doctest

samanklesaria committed 32 days ago

e2d7ba2d

Automated Code Change

h-joo committed 36 days ago

18f8f311

Add RoPE embeddings

samanklesaria committed 38 days ago

24f467da

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

a-googler committed 39 days ago

b196b016

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

dparikh79 committed 39 days ago

b1b1c184

Support dynamic attributes for Pyrefly in classes using HAS_DYNAMIC_ATTRIBUTES

QEDady committed 40 days ago

07560a7a

merge main into validate-dropout-rate

mohsinm-dev committed 41 days ago

Verified 8f75c1cc

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

a-googler committed 42 days ago

29e55c36

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

a-googler committed 42 days ago

51b3d8f1

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

a-googler committed 42 days ago

067a571d

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

a-googler committed 42 days ago

dd946850

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

chenkuanliao committed 44 days ago

ca14973a

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

a-googler committed 46 days ago

dcfabd05

Update tensorflow dependency versions in pyproject.toml

vfdev-5 committed 46 days ago

Verified 395e0331

update to v0.12.8

Cristian Garcia committed 49 days ago

3a491108

validate dropout_rate in dot_product_attention

mohsinm-dev committed 50 days ago

Verified 89b5e6b8

optimize jit_partial

Cristian Garcia committed 51 days ago

0a0aca74

Fixed typos in docs_nnx/migrating/pytorch_to_jax_flax.rst

vfdev-5 committed 51 days ago

b070daad

replace bare asserts with ValueError in attention validation

mohsinm-dev committed 51 days ago

4f85c88e

Older