google/flax

Pull Requests Commits

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

a-googler committed 1 year ago

15795867

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

a-googler committed 1 year ago

bc93cb1d

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

a-googler committed 1 year ago

ea7f9df1

[nnx] improve toy examples

cgarciae committed 1 year ago

6d638e60

[nnx] add Dropout.rngs

cgarciae committed 1 year ago

1a1ca486

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

a-googler committed 1 year ago

2b257b45

[nnx] expose Sequential

cgarciae committed 1 year ago

2821c64b

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

a-googler committed 1 year ago

e4bfdaf5

[nnx] fix nnx basics

cgarciae committed 1 year ago

fbfc0d0b

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

a-googler committed 1 year ago

22ed17e1

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

a-googler committed 1 year ago

f82b1535

[nnx] grad accepts argnums

cgarciae committed 1 year ago

febc1c7a

[nnx] RngKey tags

cgarciae committed 1 year ago

c53c17d7

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

a-googler committed 1 year ago

30fb7ff9

Merge pull request #3782 from chiamp:nnx_module_subclass

a-googler committed 1 year ago

5d9fd9fb

[nnx] allow all jit kwargs in nnx.jit

cgarciae committed 1 year ago

cfec9bb1

have TrainState subclass nnx.Module

chiamp committed 1 year ago

543686b1

Remove the uninitialized field "embedding" in nn.Embed

kho committed 1 year ago

f2bdcd82

[nnx] remove copy in graph unflatten

cgarciae committed 1 year ago

f7ab4070

[nnx] non-str State keys

cgarciae committed 1 year ago

1d5f7185

[nnx] simplify readme

cgarciae committed 1 year ago

9788dd46

Merge pull request #3794 from google:nnx-static-goes-first

a-googler committed 1 year ago

cc740d4d

[nnx] split returns graphdef first

cgarciae committed 1 year ago

1739c522

Merge pull request #3788 from google:nnx-train-state-struct

a-googler committed 1 year ago

9c4c8f58

Merge pull request #3793 from google:nnx-rngs-are-nodes

a-googler committed 1 year ago

f4337c3c

[nnx] Rngs and RngStream inherit from GraphNode

cgarciae committed 1 year ago

453f38bf

[nnx] TrainState uses struct

cgarciae committed 1 year ago

51b97edd

Merge pull request #3796 from chiamp:graphnode

a-googler committed 1 year ago

514c1119

expose nnx.GraphNode

chiamp committed 1 year ago

1eb988d8

Exposing the experimental _split_transpose JAX scan parameter in Flax.

a-googler committed 1 year ago

4671ddd9

Newer Older