google/flax

Pull Requests Commits

[linen] Module split/merge

cgarciae committed 1 year ago

df736488

Remove some outdated backward-compatibility code.

IvyZX committed 1 year ago

5c971436

Modify Flax checkpointing in preparation for cl/650338576.

cpgaffney1 committed 1 year ago

e59b14e9

Merge pull request #4070 from google:nnx-improve-mnist-tutorial

a-googler committed 1 year ago

1b58348d

Merge pull request #4066 from IvyZX:bridge

a-googler committed 1 year ago

359ec7fd

[nnx] improve mnist tutorial

cgarciae committed 1 year ago

ebfe1152

Rename compat to bridge

IvyZX committed 1 year ago

2389c397

Merge pull request #4067 from IvyZX:misc

a-googler committed 1 year ago

0c3c74c8

uncomment some tests after 3.9 dropped

IvyZX committed 1 year ago

e800d5b0

Merge pull request #4036 from kaushaladiti-2802:patch-1

a-googler committed 1 year ago

36612a38

Update orbax_upgrade_guide.rst

kaushaladiti-2802 committed 1 year ago

Verified 49100052

Merge pull request #4065 from google:fix-accuracy

a-googler committed 1 year ago

0da4f31d

Update orbax_upgrade_guide.rst

kaushaladiti-2802 committed 1 year ago

Verified 7298cc5c

[nnx] fix Accuracy on eager mode

cgarciae committed 1 year ago

26d5e68d

Update orbax_upgrade_guide.rst

kaushaladiti-2802 committed 1 year ago

Verified 5406692c

Update orbax_upgrade_guide.rst

kaushaladiti-2802 committed 1 year ago

Verified 8e307f10

Merge pull request #4063 from google:fix-notebook-links

a-googler committed 1 year ago

4e911139

Merge pull request #3984 from kaixih:rename_fp8_custom_dtype

a-googler committed 1 year ago

b3785dc7

[nnx] fix mnist_tutorial colab link

cgarciae committed 1 year ago

0be01f9e

Stop writing msgpack file for new checkpoints and update empty nodes handling so that it no longer depends on this file.

dubey committed 1 year ago

96cf6490

Merge pull request #4030 from google:nnx-pytree

a-googler committed 1 year ago

96bb44ca

[nnx] stabilize unsafe_pytree

cgarciae committed 1 year ago

4ed137c5

Stop writing msgpack file for new checkpoints and update empty nodes handling so that it no longer depends on this file.

cpgaffney1 committed 1 year ago

0fb17771

Update orbax_upgrade_guide.rst for async checkpointing usage examples

kaushaladiti-2802 committed 1 year ago

4b4d2fee

Merge pull request #4049 from google:nnx-fix-variable-inplace-operators

a-googler committed 1 year ago

1367acb1

kaixih committed 1 year ago

798cfe7b

Rename the new dtype

kaixih committed 1 year ago

16f28666

Merge pull request #4047 from rajasekharporeddy:main

a-googler committed 1 year ago

3112da6f

Add functool.wraps() annotation to flax.nn.jit.

hawkinsp committed 1 year ago

37123d52

[nnx] fix Variable overloads and add shape/dtype properties

cgarciae committed 1 year ago

385ca738

Older