google/flax

Pull Requests Commits

[pmap] In-line definitions of `jax.device_put_sharded` and `jax.device_put_replicated`.

danielsuo committed 4 days ago

640ecfe6

Merge pull request #5353 from samanklesaria:fix_orbax_export_dep

a-googler committed 4 days ago

35d911a2

Merge pull request #5231 from samanklesaria:opt_sharding

a-googler committed 4 days ago

bb50f212

fix transform_metadata

Cristian Garcia committed 8 days ago

a138d9f3

add nnx.map and nnx.abstract_with_sharding

Cristian Garcia committed 8 days ago

1572c84e

Merge pull request #5246 from samanklesaria:generalize_out_sharding

a-googler committed 8 days ago

3bcc4a91

improve aliasing error msg

Cristian Garcia committed 8 days ago

1a567002

add graph node in prefix checks

Cristian Garcia committed 8 days ago

9eb565cd

improve nnx.Dict error handling

Cristian Garcia committed 8 days ago

ebc81850

simplify SimpleScan

Cristian Garcia committed 8 days ago

2965f95a

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

yashk2810 committed 9 days ago

f521e955

Do a few more cleanups after pmap_shmap merge

hawkinsp committed 9 days ago

50433ef0

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 10 days ago

2b3ffaf3

Use with_sharding_constraint for Format application

samanklesaria committed 10 days ago

1c2cec46

Generalize out_sharding to work with NamedSharding and Format

samanklesaria committed 10 days ago

0399b9d3

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

Move orbax-export to optional deps

samanklesaria committed 11 days ago

c4d15662

check aliases on all transform args and simplify apply_variable_updates

Cristian Garcia committed 11 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

Older