google/flax

Pull Requests Commits

[nnx] improve nnx basics

cgarciae committed 1 year ago

1f27faa8

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

7d3ad179

[linen] fix DenseGeneral init

cgarciae committed 2 years ago

36ee794e

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

a-googler committed 2 years ago

bf308cad

[nnx] fix nnx_basics

cgarciae committed 2 years ago

7566585c

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

a-googler committed 2 years ago

f161a4f7

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

a-googler committed 2 years ago

1f39ba26

[nnx] remove pytreelib

cgarciae committed 2 years ago

69519f28

[linen] fix cloudpickle sentinel cloning

cgarciae committed 2 years ago

d00dfc7f

Merge pull request #3792 from chiamp:autosummary

a-googler committed 2 years ago

5ce3485c

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

a-googler committed 2 years ago

27184559

fixed autosummary

chiamp committed 2 years ago

1f6cd8a3

Merge pull request #3810 from chiamp:nnx_guide

a-googler committed 2 years ago

91cfc2a4

Merge pull request #3819 from chiamp:jit

a-googler committed 2 years ago

39ce98fe

Merge pull request #3823 from chiamp:tree_map

a-googler committed 2 years ago

054b5eb1

updated nnx mnist tutorial

chiamp committed 2 years ago

823c6160

added nnx vs jax transforms guide

chiamp committed 2 years ago

d87075fb

updated tree_map

chiamp committed 2 years ago

972c5521

Merge pull request #3828 from chiamp:filter

a-googler committed 2 years ago

5199d5cf

added tree_map deprecation warning filter

chiamp committed 2 years ago

04313c08

Merge pull request #3818 from chiamp:optimizer

a-googler committed 2 years ago

6b03439c

Merge pull request #3776 from chiamp:dropout_broadcast

a-googler committed 2 years ago

3a6c399e

Merge pull request #3789 from chiamp:optax

a-googler committed 2 years ago

7b47a753

fixed optax guide links and docstring typos

chiamp committed 2 years ago

a375256d

added dropout broadcast test

chiamp committed 2 years ago

46a53f22

relaxed kwarg for

chiamp committed 2 years ago

bb959624

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

a-googler committed 2 years ago

760cf868

Merge pull request #3807 from google:nnx-rng-tags

a-googler committed 2 years ago

15795867

Older