google/flax

Pull Requests Commits

[nnx] minor changes to mnist tutorial

cgarciae committed 2 years ago

1be6a6c0

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

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] non-str State keys

cgarciae committed 2 years ago

1d5f7185

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

Exposing the experimental _split_transpose JAX scan parameter in Flax.

a-googler committed 2 years ago

4671ddd9

Merge pull request #3783 from google:nnx-jit-accepts-many-Modules

a-googler committed 2 years ago

fe995740

[nnx] jit accepts many Modules

cgarciae committed 2 years ago

e452f5a4

Merge pull request #3790 from google:nnx-graph-node-base

a-googler committed 2 years ago

6f1f1ef0

[nnx] add GraphNode base class

cgarciae committed 2 years ago

6948138a

Merge pull request #3791 from google:nnx-arrays-are-state

a-googler committed 2 years ago

0ab03658

[nnx] Arrays are state

cgarciae committed 2 years ago

87181104

Merge pull request #3785 from chiamp:pytreenode

a-googler committed 2 years ago

b57e3aa7

extended pytreenode to take dataclass kwargs

chiamp committed 2 years ago

9ccd27d5

Merge pull request #3651 from chiamp:kw_only

a-googler committed 2 years ago

135b6527

Merge pull request #3772 from chiamp:tabulate_norm

a-googler committed 2 years ago

c53a5151

added kw_only dataclass test

chiamp committed 2 years ago

78e8056d

Use short-circuiting access to debug_key_reuse

Jake VanderPlas committed 2 years ago

fb3107fd

Account for JAX flag name change in https://github.com/google/jax/pull/20370

Jake VanderPlas committed 2 years ago

f234fc1b

Merge pull request #3768 from google:nnx-pytree-are-trees

a-googler committed 2 years ago

12b919c9

Merge pull request #3773 from chiamp:mgu_test

a-googler committed 2 years ago

20213c97

Older