google/flax

Pull Requests Commits

Remove spmd_lowering from MeshComputation and do more cleanups after pmap, PmapSharding is gone

hawkinsp committed 9 days ago

fde3b2e8

update to version 0.12.6

Cristian Garcia committed 9 days ago

ee6d3c58

Do a few more cleanups after pmap_shmap merge

yashk2810 committed 9 days ago

b96a07a0

Merge pull request #5332 from google:metrics-mask

a-googler committed 9 days ago

2b3ffaf3

fix jit_partial lower and compile

Cristian Garcia committed 10 days ago

15683fd1

Merge pull request #5257 from samanklesaria:capture_intms

a-googler committed 10 days ago

e5a6ff6d

check aliases on all transform args and simplify apply_variable_updates

Cristian Garcia committed 10 days ago

1fcc2a5c

remove aliases from nnx.graph in favor of nnx.compat

Cristian Garcia committed 11 days ago

b313d06d

Introduce `manual_type: ManualAxisType` parameter on ShapedArray to track varying/unreduced/reduced and remove `vma: frozenset` parameter.

yashk2810 committed 11 days ago

4ba66fce

improve error messages for tree mode duplicates check

Cristian Garcia committed 11 days ago

d930ccd3

clean up custom_vjp's graph_updates=False section

Cristian Garcia committed 11 days ago

93c430db

add transform_metadata transform

Cristian Garcia committed 11 days ago

1779706a

allow in/out_axes when graph_updates=False

Cristian Garcia committed 11 days ago

4f58b353

Add nnx.capture

samanklesaria committed 11 days ago

0e926241

support hijax and ref Variables in simple transforms

Cristian Garcia committed 11 days ago

a9059fc2

don't allow Variable mutation in custom_vjp on differentiable arguments

Cristian Garcia committed 11 days ago

d812e604

add more tests that check for consistent aliasing in transforms

Cristian Garcia committed 11 days ago

df5c1801

add compat module

Cristian Garcia committed 12 days ago

5a79a960

add graph_updates to eval_shape and checkify

Cristian Garcia committed 12 days ago

6b358ed7

add graph_updates to vjp and jvp

Cristian Garcia committed 12 days ago

4386571e

add graph_updates to remat

Cristian Garcia committed 12 days ago

47070f22

add graph_updates for custom_vjp

Cristian Garcia committed 12 days ago

1e6e34bf

Added support for data masking in Average, Accuracy and MultiMetric

vfdev-5 committed 13 days ago

5c7fb394

Copybara import of the project:

Cristian Garcia committed 15 days ago

4183e15c

add tree-mode-nnx FLIP

Cristian Garcia committed 15 days ago

74e93035

add graph_updates to pmap, grad, and value_and_grad

Cristian Garcia committed 15 days ago

90df65b8

add graph_updates to fori_loop

Cristian Garcia committed 15 days ago

a2aa0ebe

Merge pull request #5326 from google:metrics-mask

a-googler committed 15 days ago

4cf2b8c5

add graph_updates to while_loop

Cristian Garcia committed 16 days ago

a8ecd069

add graph_updates for scan

Cristian Garcia committed 16 days ago

56471a47

Older