google/flax

Pull Requests Commits

Add graph=False support for nnx.grad and nnx.value_and_grad

Cristian Garcia committed 1 day ago

851b585c

Add tree-mode support for nnx.shard_map

Cristian Garcia committed 1 day ago

0bfe50a5

Add tree-mode jit (graph=False) to nnx.jit

Cristian Garcia committed 1 day ago

b528725b

Cristian Garcia committed 1 day ago

ae0ed509

Adds a test for `linen.WeightNorm` and `linen.with_partitioning` compatibility

jimlinntu committed 2 days ago

cf350e9a

No public description

a-googler committed 2 days ago

10c1564c

maintain data/static definition in split / state

Cristian Garcia committed 2 days ago

404a4208

Makes `linen.Partitioned` compatible with `linen.WeightNorm`

jimlinntu committed 2 days ago

7814e6c5

Merge pull request #5230 from gurusai-voleti:ai-gsutil-migration-aa63244b02d84e3bbc066fe1da755c37

a-googler committed 3 days ago

5b1ebbd6

Merge pull request #5180 from ayulockedin:feat/nnx-gqa-support

a-googler committed 3 days ago

a972a4df

Merge pull request #5218 from samanklesaria:issues/5217

a-googler committed 3 days ago

afea26fa

Add tests for Sequential static leaves and map

samanklesaria committed 4 days ago

4fa02af4

chore: Migrate gsutil usage to gcloud storage

gurusai-voleti committed 4 days ago

5d7e13a3

Merge pull request #5222 from samanklesaria:doctest_check

a-googler committed 6 days ago

d87f7da7

Merge pull request #5225 from samanklesaria:commit_to_test_ci

a-googler committed 6 days ago

79992cc9

Improve Pytree flatten/unflatten

Cristian Garcia committed 6 days ago

2ab6208e

Pin optax dependency to avoid error

samanklesaria committed 6 days ago

3c2277bf

Merge pull request #5205 from samanklesaria:embed_out_sharding

a-googler committed 8 days ago

2e1f5adc

Merge pull request #5220 from chapman20j:bonsai_link

a-googler committed 8 days ago

93d00808

Pin tqdm version for doctest failures

samanklesaria committed 9 days ago

5bc09657

chapman20j committed 9 days ago

172bafe5

Merge pull request #5204 from chapman20j:view

a-googler committed 10 days ago

4089582f

Merge pull request #5212 from google:fix-dropout-5209

a-googler committed 10 days ago

be338919

chapman20j committed 13 days ago

4f5a64c6

Copybara import of the project:

Cristian Garcia committed 13 days ago

b4e73d79

Fixed mask sharding if inputs is sharded

vfdev-5 committed 15 days ago

0bff68cf

Copybara import of the project:

Cristian Garcia committed 15 days ago

df2ed160

[pmap] Speed up the new `jax.pmap` implementation.

danielsuo committed 16 days ago

01c0bd13

Merge pull request #5072 from google:fix-list

a-googler committed 16 days ago

5adb1744

fix List and improve Pytree

cgarciae committed 16 days ago

af408cf5

Older