google/flax

Pull Requests Commits

Rename `SDS(vma: frozenset = ...)` to `SDS(manual_type: jax.sharding.ManualAxisType = ...)`

yashk2810 committed 25 days ago

0dd6cb22

clean up custom_vjp's graph_updates=False section

Cristian Garcia committed 25 days ago

93c430db

add transform_metadata transform

Cristian Garcia committed 25 days ago

1779706a

allow in/out_axes when graph_updates=False

Cristian Garcia committed 25 days ago

4f58b353

support hijax and ref Variables in simple transforms

Cristian Garcia committed 25 days ago

a9059fc2

don't allow Variable mutation in custom_vjp on differentiable arguments

Cristian Garcia committed 25 days ago

d812e604

add more tests that check for consistent aliasing in transforms

Cristian Garcia committed 25 days ago

df5c1801

add compat module

Cristian Garcia committed 26 days ago

5a79a960

add graph_updates to eval_shape and checkify

Cristian Garcia committed 26 days ago

6b358ed7

add graph_updates to vjp and jvp

Cristian Garcia committed 26 days ago

4386571e

add graph_updates to remat

Cristian Garcia committed 26 days ago

47070f22

add graph_updates for custom_vjp

Cristian Garcia committed 26 days ago

1e6e34bf

Copybara import of the project:

Cristian Garcia committed 29 days ago

4183e15c

add tree-mode-nnx FLIP

Cristian Garcia committed 29 days ago

74e93035

add graph_updates to pmap, grad, and value_and_grad

Cristian Garcia committed 29 days ago

90df65b8

add graph_updates to fori_loop

Cristian Garcia committed 29 days ago

a2aa0ebe

Merge pull request #5326 from google:metrics-mask

a-googler committed 29 days ago

4cf2b8c5

add graph_updates to while_loop

Cristian Garcia committed 30 days ago

a8ecd069

add graph_updates for scan

Cristian Garcia committed 30 days ago

56471a47

Merge pull request #5270 from samanklesaria:split_method

a-googler committed 30 days ago

d4653e55

add graph_updates to vmap

Cristian Garcia committed 30 days ago

35972c24

Remove jax/tools/colab_tpu.py.

hawkinsp committed 30 days ago

8ed67b85

Added support for data masking in Average, Accuracy and MultiMetric

vfdev-5 committed 30 days ago

5074fb68

add graph_updates argument in shard_map

Cristian Garcia committed 30 days ago

115dcf58

Add split method to RngStream

samanklesaria committed 30 days ago

696328ab

add graph_updates argument to jit

Cristian Garcia committed 31 days ago

a1899fd4

Allow DenyList to be compared.

levskaya committed 33 days ago

22468a1b

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

a-googler committed 33 days ago

9b399162

support out_sharding mapping in set_metadata

Cristian Garcia committed 36 days ago

4614ac2b

fix eval_shape's _to_variable

Cristian Garcia committed 36 days ago

8335dfe1

Older