google/flax

Pull Requests Commits

Upgrade third_party/python_runtime/v3_13_unstable to 3.13.11

a-googler committed 15 hours ago

8955e87f

fix transform_metadata

Cristian Garcia committed 3 days ago

a138d9f3

add nnx.map and nnx.abstract_with_sharding

Cristian Garcia committed 4 days ago

1572c84e

Merge pull request #5246 from samanklesaria:generalize_out_sharding

a-googler committed 4 days ago

3bcc4a91

improve aliasing error msg

Cristian Garcia committed 4 days ago

1a567002

add graph node in prefix checks

Cristian Garcia committed 4 days ago

9eb565cd

improve nnx.Dict error handling

Cristian Garcia committed 4 days ago

ebc81850

simplify SimpleScan

Cristian Garcia committed 4 days ago

2965f95a

Do a few more cleanups after the pmap shmap merge deletion.

yashk2810 committed 4 days ago

f521e955

Do a few more cleanups after pmap_shmap merge

hawkinsp committed 4 days ago

50433ef0

update to version 0.12.6

Cristian Garcia committed 4 days ago

ee6d3c58

Do a few more cleanups after pmap_shmap merge

yashk2810 committed 5 days ago

b96a07a0

Merge pull request #5332 from google:metrics-mask

a-googler committed 5 days ago

2b3ffaf3

Use with_sharding_constraint for Format application

samanklesaria committed 5 days ago

1c2cec46

Generalize out_sharding to work with NamedSharding and Format

samanklesaria committed 5 days ago

0399b9d3

fix jit_partial lower and compile

Cristian Garcia committed 6 days ago

15683fd1

Merge pull request #5257 from samanklesaria:capture_intms

a-googler committed 6 days ago

e5a6ff6d

check aliases on all transform args and simplify apply_variable_updates

Cristian Garcia committed 6 days ago

1fcc2a5c

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

Cristian Garcia committed 6 days ago

b313d06d

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

yashk2810 committed 6 days ago

4ba66fce

improve error messages for tree mode duplicates check

Cristian Garcia committed 6 days ago

d930ccd3

clean up custom_vjp's graph_updates=False section

Cristian Garcia committed 6 days ago

93c430db

add transform_metadata transform

Cristian Garcia committed 7 days ago

1779706a

allow in/out_axes when graph_updates=False

Cristian Garcia committed 7 days ago

4f58b353

Add nnx.capture

samanklesaria committed 7 days ago

0e926241

support hijax and ref Variables in simple transforms

Cristian Garcia committed 7 days ago

a9059fc2

don't allow Variable mutation in custom_vjp on differentiable arguments

Cristian Garcia committed 7 days ago

d812e604

add more tests that check for consistent aliasing in transforms

Cristian Garcia committed 7 days ago

df5c1801

add compat module

Cristian Garcia committed 8 days ago

5a79a960

add graph_updates to eval_shape and checkify

Cristian Garcia committed 8 days ago

6b358ed7

Older