google/flax

Pull Requests Commits

Add nnx_graph_mode config flag

Cristian Garcia committed 50 days ago

ea56e5e7

Add graph argument to split, state, graphdef, clone

Cristian Garcia committed 50 days ago

fdf19052

Improve unregistered data detection

Cristian Garcia committed 50 days ago

431291ed

[flax:benchmarks] Fix flax benchmarks.

danielsuo committed 51 days ago

4a15be78

Implemented f'string

ayulockedin committed 52 days ago

91fea70e

feat(nnx): add GQA support to MultiHeadAttention

ayulockedin committed 52 days ago

23918aea

Add tree-mode arg error handling

Cristian Garcia committed 52 days ago

255eea65

Add tree-mode support for nnx.vmap

Cristian Garcia committed 56 days ago

3d634412

Add tree-mode support for nnx.remat

Cristian Garcia committed 56 days ago

aed27f67

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

Cristian Garcia committed 56 days ago

49213e07

maintain data/static definition in split / state

Cristian Garcia committed 57 days ago

7811accc

Rename sharding_names to out_sharding in NNX Variable metadata

Cristian Garcia committed 58 days ago

a5c76831

Add tree-mode support for nnx.shard_map

Cristian Garcia committed 58 days ago

0bfe50a5

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

Cristian Garcia committed 58 days ago

b528725b

Cristian Garcia committed 58 days ago

ae0ed509

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

jimlinntu committed 59 days ago

cf350e9a

No public description

a-googler committed 59 days ago

10c1564c

maintain data/static definition in split / state

Cristian Garcia committed 60 days ago

404a4208

Fix PyTreeNode + Generic losing __parameters__ when Generic is last in bases (#5233)

mohsinm-dev committed 60 days ago

Verified 4dd61948

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

jimlinntu committed 60 days ago

7814e6c5

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

a-googler committed 60 days ago

5b1ebbd6

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

a-googler committed 60 days ago

a972a4df

Merge pull request #5218 from samanklesaria:issues/5217

a-googler committed 60 days ago

afea26fa

Add tests for Sequential static leaves and map

samanklesaria committed 61 days ago

4fa02af4

chore: Migrate gsutil usage to gcloud storage

gurusai-voleti committed 61 days ago

5d7e13a3

Removed redundant code

vfdev-5 committed 61 days ago

985141fd

Merge pull request #5222 from samanklesaria:doctest_check

a-googler committed 63 days ago

d87f7da7

Merge pull request #5225 from samanklesaria:commit_to_test_ci

a-googler committed 63 days ago

79992cc9

Improve Pytree flatten/unflatten

Cristian Garcia committed 64 days ago

2ab6208e

Pin optax dependency to avoid error

samanklesaria committed 64 days ago

3c2277bf

Newer Older