google/flax

Pull Requests Commits

Merge pull request #3817 from google:nnx-jit-apply-sharding

a-googler committed 1 year ago

2438518b

[nnx] jit constrain object state

cgarciae committed 1 year ago

8192b966

Merge pull request #3834 from google:linen-fix-dense-general-init

a-googler committed 1 year ago

7d3ad179

[linen] fix DenseGeneral init

cgarciae committed 1 year ago

36ee794e

Add example of unbox() and replace_boxed()

IvyZX committed 1 year ago

159dfd43

Merge pull request #3839 from google:nnx-fix-nnx-basics

a-googler committed 1 year ago

bf308cad

[nnx] fix nnx_basics

cgarciae committed 1 year ago

7566585c

Merge pull request #3816 from google:nnx-remove-treelib

a-googler committed 1 year ago

f161a4f7

Merge pull request #3825 from google:linen-cloudpickle-issue

a-googler committed 1 year ago

1f39ba26

[nnx] remove pytreelib

cgarciae committed 1 year ago

69519f28

[linen] fix cloudpickle sentinel cloning

cgarciae committed 1 year ago

d00dfc7f

Merge pull request #3792 from chiamp:autosummary

a-googler committed 1 year ago

5ce3485c

Merge pull request #3815 from google:nnx-dropout-optional-state

a-googler committed 1 year ago

27184559

fixed autosummary

chiamp committed 1 year ago

1f6cd8a3

Merge pull request #3810 from chiamp:nnx_guide

a-googler committed 1 year ago

91cfc2a4

Update the type rule to support non-jit use cases

kaixih committed 1 year ago

fd9c467f

Merge pull request #3819 from chiamp:jit

a-googler committed 1 year ago

39ce98fe

Merge pull request #3823 from chiamp:tree_map

a-googler committed 1 year ago

054b5eb1

updated nnx mnist tutorial

chiamp committed 1 year ago

823c6160

added nnx vs jax transforms guide

chiamp committed 1 year ago

d87075fb

updated tree_map

chiamp committed 1 year ago

972c5521

Merge pull request #3828 from chiamp:filter

a-googler committed 1 year ago

5199d5cf

added tree_map deprecation warning filter

chiamp committed 1 year ago

04313c08

Merge pull request #3818 from chiamp:optimizer

a-googler committed 1 year ago

6b03439c

Merge pull request #3776 from chiamp:dropout_broadcast

a-googler committed 1 year ago

3a6c399e

Merge pull request #3789 from chiamp:optax

a-googler committed 1 year ago

7b47a753

fixed optax guide links and docstring typos

chiamp committed 1 year ago

a375256d

added dropout broadcast test

chiamp committed 1 year ago

46a53f22

relaxed kwarg for

chiamp committed 1 year ago

bb959624

Merge pull request #3804 from google:nnx-unflatten-no-copy

a-googler committed 1 year ago

760cf868

Newer Older