google/flax

Pull Requests Commits

[nnx] experimental transforms

cgarciae committed 2 years ago

a19dbf04

Merge pull request #4078 from google:nnx-extract-apis

a-googler committed 2 years ago

e7f43dd7

[nnx] add extract APIs

cgarciae committed 2 years ago

16c2c324

Merge pull request #4076 from kaixih:add_fp8_doc

a-googler committed 2 years ago

e8021e9f

Add doc for fp8

kaixih committed 2 years ago

350d9f83

Reverting unsafe_pytree #4030

Cristian Garcia committed 2 years ago

239d4e6b

Remove some outdated backward-compatibility code.

IvyZX committed 2 years ago

5c971436

Modify Flax checkpointing in preparation for cl/650338576.

cpgaffney1 committed 2 years ago

e59b14e9

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

a-googler committed 2 years ago

1b58348d

Merge pull request #4066 from IvyZX:bridge

a-googler committed 2 years ago

359ec7fd

[nnx] improve mnist tutorial

cgarciae committed 2 years ago

ebfe1152

Rename compat to bridge

IvyZX committed 2 years ago

2389c397

Merge pull request #4067 from IvyZX:misc

a-googler committed 2 years ago

0c3c74c8

uncomment some tests after 3.9 dropped

IvyZX committed 2 years ago

e800d5b0

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

a-googler committed 2 years ago

36612a38

Update orbax_upgrade_guide.rst

kaushaladiti-2802 committed 2 years ago

Verified 49100052

Merge pull request #4065 from google:fix-accuracy

a-googler committed 2 years ago

0da4f31d

Update orbax_upgrade_guide.rst

kaushaladiti-2802 committed 2 years ago

Verified 7298cc5c

[nnx] fix Accuracy on eager mode

cgarciae committed 2 years ago

26d5e68d

Update orbax_upgrade_guide.rst

kaushaladiti-2802 committed 2 years ago

Verified 5406692c

Update orbax_upgrade_guide.rst

kaushaladiti-2802 committed 2 years ago

Verified 8e307f10

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

a-googler committed 2 years ago

4e911139

Merge pull request #3984 from kaixih:rename_fp8_custom_dtype

a-googler committed 2 years ago

b3785dc7

[nnx] fix mnist_tutorial colab link

cgarciae committed 2 years 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 2 years ago

96cf6490

Merge pull request #4030 from google:nnx-pytree

a-googler committed 2 years ago

96bb44ca

[nnx] stabilize unsafe_pytree

cgarciae committed 2 years 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 2 years ago

0fb17771

Update orbax_upgrade_guide.rst for async checkpointing usage examples

kaushaladiti-2802 committed 2 years ago

4b4d2fee

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

a-googler committed 2 years ago

1367acb1

Older