google/flax

Pull Requests Commits

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

[nnx] add Dropout.rngs

cgarciae committed 2 years ago

1a1ca486

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

Remove the uninitialized field "embedding" in nn.Embed

kho committed 2 years ago

f2bdcd82

[nnx] remove copy in graph unflatten

cgarciae committed 2 years ago

f7ab4070

[nnx] non-str State keys

cgarciae committed 2 years ago

1d5f7185

[nnx] simplify readme

cgarciae committed 2 years ago

9788dd46

Newer Older