google/flax

Pull Requests Commits

Refactor NNX Variable to remove legacy JAX QDD dependencies.

Cristian Garcia committed 5 days ago

c31deb2b

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

a-googler committed 7 days ago

ae2d29e4

implement axes_scan's broadcast semantics via dce

mattjj committed 7 days ago

ba905581

Merge pull request #5481 from samanklesaria:rope

a-googler committed 13 days ago

a5252645

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

a-googler committed 13 days ago

7537c605

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

a-googler committed 15 days ago

92cf9a7d

simplify vmapped metrics reset test

chenkuanliao committed 28 days ago

abc2dbe7

Merge pull request #5526 from samanklesaria:fix_ci_precision

a-googler committed 29 days ago

a31129e2

Use less precise numbers in linen doctest

samanklesaria committed 30 days ago

e2d7ba2d

Automated Code Change

h-joo committed 34 days ago

18f8f311

Add RoPE embeddings

samanklesaria committed 36 days ago

24f467da

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

a-googler committed 37 days ago

b196b016

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

dparikh79 committed 37 days ago

b1b1c184

Support dynamic attributes for Pyrefly in classes using HAS_DYNAMIC_ATTRIBUTES

QEDady committed 38 days ago

07560a7a

merge main into validate-dropout-rate

mohsinm-dev committed 40 days ago

Verified 8f75c1cc

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

a-googler committed 40 days ago

29e55c36

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

a-googler committed 40 days ago

51b3d8f1

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

a-googler committed 40 days ago

067a571d

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

a-googler committed 40 days ago

dd946850

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

chenkuanliao committed 42 days ago

ca14973a

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

a-googler committed 44 days ago

dcfabd05

Update tensorflow dependency versions in pyproject.toml

vfdev-5 committed 44 days ago

Verified 395e0331

update to v0.12.8

Cristian Garcia committed 47 days ago

3a491108

validate dropout_rate in dot_product_attention

mohsinm-dev committed 49 days ago

Verified 89b5e6b8

optimize jit_partial

Cristian Garcia committed 49 days ago

0a0aca74

Fixed typos in docs_nnx/migrating/pytorch_to_jax_flax.rst

vfdev-5 committed 49 days ago

b070daad

replace bare asserts with ValueError in attention validation

mohsinm-dev committed 49 days ago

4f85c88e

Added Object Detection example with minimal DETR

vfdev-5 committed 49 days ago

25785b10

fix nnx promote_dtype docstring

mohsinm-dev committed 49 days ago

1f3d9549

Merge pull request #5507 from google:add-dependabot

a-googler committed 49 days ago

0682fe4d

Older