google/flax

Pull Requests Commits

[nnx] refactor graph_utils

cgarciae committed 2 years ago

0e8f66c9

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

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

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

a-googler committed 2 years ago

cc740d4d

[nnx] split returns graphdef first

cgarciae committed 2 years ago

1739c522

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

a-googler committed 2 years ago

9c4c8f58

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

a-googler committed 2 years ago

f4337c3c

[nnx] Rngs and RngStream inherit from GraphNode

cgarciae committed 2 years ago

453f38bf

[nnx] TrainState uses struct

cgarciae committed 2 years ago

51b97edd

Merge pull request #3796 from chiamp:graphnode

a-googler committed 2 years ago

514c1119

expose nnx.GraphNode

chiamp committed 2 years ago

1eb988d8

Older