google/flax

Pull Requests Commits

Add **init_kwargs to self.param type annotation.

a-googler committed 2 years ago

c54d9b99

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

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

Merge pull request #3814 from google:nnx-expose-sequential

a-googler committed 2 years ago

bc93cb1d

Merge pull request #3813 from google:nnx-improve-lifted-example

a-googler committed 2 years ago

ea7f9df1

[nnx] improve toy examples

cgarciae committed 2 years ago

6d638e60

Merge pull request #3798 from google:nnx-improve-grad

a-googler committed 2 years ago

2b257b45

[nnx] expose Sequential

cgarciae committed 2 years ago

2821c64b

Merge pull request #3812 from google:nnx-minor-fix-nnx-basics

a-googler committed 2 years ago

e4bfdaf5

[nnx] fix nnx basics

cgarciae committed 2 years ago

fbfc0d0b

Merge pull request #3805 from google:nnx-simplify-readme-again-2

a-googler committed 2 years ago

22ed17e1

Merge pull request #3809 from google:nnx-all-jit-kwargs

a-googler committed 2 years ago

f82b1535

[nnx] grad accepts argnums

cgarciae committed 2 years ago

febc1c7a

[nnx] RngKey tags

cgarciae committed 2 years ago

c53c17d7

Merge pull request #3802 from google:nnx-graph-non-str-keys

a-googler committed 2 years ago

30fb7ff9

Merge pull request #3782 from chiamp:nnx_module_subclass

a-googler committed 2 years ago

5d9fd9fb

[nnx] allow all jit kwargs in nnx.jit

cgarciae committed 2 years ago

cfec9bb1

have TrainState subclass nnx.Module

chiamp committed 2 years ago

543686b1

Older