google/flax

Pull Requests Commits

[nnx] use functional api

cgarciae committed 1 year ago

27a053d9

[nnx] default runtime attributes + train/eval methods

cgarciae committed 1 year ago

1dea4990

cgarciae committed 1 year ago

5f93e836

Copybara import of the project:

Cristian Garcia committed 1 year ago

1fbe17d8

Merge pull request #3861 from gnecula:nnx_doc1

a-googler committed 1 year ago

99736930

integrate changes to nnx_basics

cgarciae committed 1 year ago

7f1b36a7

Merge pull request #3099 from google:rnncell-refactor-flip

a-googler committed 1 year ago

bf4dffff

[nnx] Some small documentation suggestions.

gnecula committed 1 year ago

05a72fc6

Merge pull request #3859 from google:nnx-improve-nnx-basics

a-googler committed 1 year ago

f7c0669d

Merge pull request #3843 from IvyZX:jit-guide

a-googler committed 1 year ago

9c008e33

[nnx] improve nnx basics

cgarciae committed 1 year ago

1f27faa8

Copybara import of the project:

Cristian Garcia committed 1 year ago

3711fa60

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

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

Older