google/flax

Pull Requests Commits

Add tree-mode support to nnx.scan

Cristian Garcia committed 16 minutes ago

934d53cd

Add tree-mode support to nnx.{cond,switch,eval_shape,checkify}

Cristian Garcia committed 46 minutes ago

761a27d3

Add nnx_graph_mode config flag

Cristian Garcia committed 2 hours ago

ea56e5e7

Add graph argument to split, state, graphdef, clone

Cristian Garcia committed 3 hours ago

fdf19052

Improve unregistered data detection

Cristian Garcia committed 3 hours ago

431291ed

[flax:benchmarks] Fix flax benchmarks.

danielsuo committed 1 day ago

4a15be78

Add tree-mode arg error handling

Cristian Garcia committed 2 days ago

255eea65

Add tree-mode support for nnx.vmap

Cristian Garcia committed 6 days ago

3d634412

Add tree-mode support for nnx.remat

Cristian Garcia committed 6 days ago

aed27f67

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

Cristian Garcia committed 6 days ago

49213e07

maintain data/static definition in split / state

Cristian Garcia committed 7 days ago

7811accc

Rename sharding_names to out_sharding in NNX Variable metadata

Cristian Garcia committed 8 days ago

a5c76831

Add tree-mode support for nnx.shard_map

Cristian Garcia committed 8 days ago

0bfe50a5

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

Cristian Garcia committed 8 days ago

b528725b

Cristian Garcia committed 8 days ago

ae0ed509

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

jimlinntu committed 10 days ago

cf350e9a

No public description

a-googler committed 10 days ago

10c1564c

maintain data/static definition in split / state

Cristian Garcia committed 10 days ago

404a4208

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

jimlinntu committed 10 days ago

7814e6c5

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

a-googler committed 10 days ago

5b1ebbd6

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

a-googler committed 11 days ago

a972a4df

Merge pull request #5218 from samanklesaria:issues/5217

a-googler committed 11 days ago

afea26fa

Add tests for Sequential static leaves and map

samanklesaria committed 11 days ago

4fa02af4

chore: Migrate gsutil usage to gcloud storage

gurusai-voleti committed 11 days ago

5d7e13a3

Merge pull request #5222 from samanklesaria:doctest_check

a-googler committed 13 days ago

d87f7da7

Merge pull request #5225 from samanklesaria:commit_to_test_ci

a-googler committed 13 days ago

79992cc9

Improve Pytree flatten/unflatten

Cristian Garcia committed 14 days ago

2ab6208e

Pin optax dependency to avoid error

samanklesaria committed 14 days ago

3c2277bf

Merge pull request #5205 from samanklesaria:embed_out_sharding

a-googler committed 15 days ago

2e1f5adc

Merge pull request #5220 from chapman20j:bonsai_link

a-googler committed 15 days ago

93d00808

Older