google/flax

Pull Requests Commits

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

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

a-googler committed 1 year ago

fe995740

[nnx] jit accepts many Modules

cgarciae committed 1 year ago

e452f5a4

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

a-googler committed 1 year ago

6f1f1ef0

[nnx] add GraphNode base class

cgarciae committed 1 year ago

6948138a

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

a-googler committed 1 year ago

0ab03658

[nnx] Arrays are state

cgarciae committed 1 year ago

87181104

Merge pull request #3785 from chiamp:pytreenode

a-googler committed 1 year ago

b57e3aa7

extended pytreenode to take dataclass kwargs

chiamp committed 1 year ago

9ccd27d5

Merge pull request #3651 from chiamp:kw_only

a-googler committed 1 year ago

135b6527

Merge pull request #3772 from chiamp:tabulate_norm

a-googler committed 1 year ago

c53a5151

added kw_only dataclass test

chiamp committed 1 year ago

78e8056d

Use short-circuiting access to debug_key_reuse

Jake VanderPlas committed 1 year ago

fb3107fd

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

Jake VanderPlas committed 1 year ago

f234fc1b

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

a-googler committed 1 year ago

12b919c9

Merge pull request #3773 from chiamp:mgu_test

a-googler committed 1 year ago

20213c97

Merge pull request #3588 from canyon289:doc_improvements

a-googler committed 1 year ago

d75c2e58

Access thread_resources via jax.interpreters.pxla instead of jax.experimental.maps

superbobry committed 1 year ago

e1465fed

added MGU reset_gate test

chiamp committed 1 year ago

cfb4d9da

fix tabulate on norm wrappers

chiamp committed 1 year ago

4f561be1

Merge pull request #3760 from carlosgmartin:reset_gate_flag

a-googler committed 1 year ago

006da41a

Add reset_gate flag to MGUCell.

carlosgmartin committed 1 year ago

f932fdb2

[nnx] pytree are trees

cgarciae committed 1 year ago

7317d9b9

Merge pull request #3769 from chiamp:note

a-googler committed 1 year ago

82201540

Merge pull request #3759 from google:nnx-improve-graph-update

a-googler committed 1 year ago

b626099b

[nnx] add ref_cache to graph_unflatten

cgarciae committed 1 year ago

c8cf3f50

use note box in docstrings

chiamp committed 1 year ago

f14d8cc9

Merge pull request #3767 from google:improve-make_rng-docstring

a-googler committed 1 year ago

79851d4f

Older