google/flax

Pull Requests Commits

add graph_updates to vjp and jvp

Cristian Garcia committed 27 days ago

4386571e

add graph_updates to remat

Cristian Garcia committed 28 days ago

47070f22

add graph_updates for custom_vjp

Cristian Garcia committed 28 days ago

1e6e34bf

Added support for data masking in Average, Accuracy and MultiMetric

vfdev-5 committed 28 days ago

5c7fb394

Copybara import of the project:

Cristian Garcia committed 30 days ago

4183e15c

add tree-mode-nnx FLIP

Cristian Garcia committed 30 days ago

74e93035

add graph_updates to pmap, grad, and value_and_grad

Cristian Garcia committed 31 days ago

90df65b8

add graph_updates to fori_loop

Cristian Garcia committed 31 days ago

a2aa0ebe

Merge pull request #5326 from google:metrics-mask

a-googler committed 31 days ago

4cf2b8c5

add graph_updates to while_loop

Cristian Garcia committed 31 days ago

a8ecd069

add graph_updates for scan

Cristian Garcia committed 31 days ago

56471a47

Merge pull request #5270 from samanklesaria:split_method

a-googler committed 32 days ago

d4653e55

add graph_updates to vmap

Cristian Garcia committed 32 days ago

35972c24

Remove jax/tools/colab_tpu.py.

hawkinsp committed 32 days ago

8ed67b85

Added support for data masking in Average, Accuracy and MultiMetric

vfdev-5 committed 32 days ago

5074fb68

add graph_updates argument in shard_map

Cristian Garcia committed 32 days ago

115dcf58

Add split method to RngStream

samanklesaria committed 32 days ago

696328ab

add graph_updates argument to jit

Cristian Garcia committed 32 days ago

a1899fd4

Allow DenyList to be compared.

levskaya committed 34 days ago

22468a1b

Merge pull request #5259 from ayulockedin:feat/gqa-support-mha

a-googler committed 35 days ago

9b399162

support out_sharding mapping in set_metadata

Cristian Garcia committed 37 days ago

4614ac2b

fix eval_shape's _to_variable

Cristian Garcia committed 38 days ago

8335dfe1

Merge pull request #5247 from samanklesaria:eval_shape_sharding

a-googler committed 38 days ago

49e34d17

Update name to optimizer_sharding

samanklesaria committed 38 days ago

61e0d32d

Resolve code review comments

samanklesaria committed 38 days ago

4e222b6a

Fix linter errors

samanklesaria committed 38 days ago

0966c0e7

Update out_sharding name

samanklesaria committed 38 days ago

Verified 8520c75b

Support explicit sharding for nnx.eval_shape

samanklesaria committed 38 days ago

6bce1438

Change 'sharding' name to 'out_sharding'

samanklesaria committed 38 days ago

3e3c4594

Added sharding propagation support in nnx.eval_shape

vfdev-5 committed 38 days ago

70351571

Newer Older