google/flax

Pull Requests Commits

have TrainState subclass nnx.Module

chiamp committed 2 years ago

297c3889

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

Merge pull request #3588 from canyon289:doc_improvements

a-googler committed 2 years ago

d75c2e58

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

superbobry committed 2 years ago

e1465fed

added MGU reset_gate test

chiamp committed 2 years ago

cfb4d9da

Merge pull request #3760 from carlosgmartin:reset_gate_flag

a-googler committed 2 years ago

006da41a

Add reset_gate flag to MGUCell.

carlosgmartin committed 2 years ago

f932fdb2

[nnx] pytree are trees

cgarciae committed 2 years ago

7317d9b9

Merge pull request #3769 from chiamp:note

a-googler committed 2 years ago

82201540

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

a-googler committed 2 years ago

b626099b

[nnx] add ref_cache to graph_unflatten

cgarciae committed 2 years ago

c8cf3f50

use note box in docstrings

chiamp committed 2 years ago

f14d8cc9

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

a-googler committed 2 years ago

79851d4f

Merge pull request #3765 from chiamp:docstring_init_apply

a-googler committed 2 years ago

2f76a92d

updated rng docstring for init and apply

chiamp committed 2 years ago

76f6c461

Merge pull request #3762 from chiamp:nnx_api_reference

a-googler committed 2 years ago

a6e771f3

Merge pull request #3744 from google:nnx-fix-demo

a-googler committed 2 years ago

e3e8cb48

Merge pull request #3758 from google:v0.8.3

a-googler committed 2 years ago

e1f672b8

use note box in make_rng docstring

cgarciae committed 2 years ago

bd831b30

added nnx api reference

chiamp committed 2 years ago

9302fc88

Make `TrainState`'s `step` possibly jax.Array. This makes `replicate` valid for type checking.

IvyZX committed 2 years ago

718aa8cc

Merge pull request #3741 from chiamp:einsum_nnx

a-googler committed 2 years ago

6844a540

Merge pull request #3743 from chiamp:nnx_value

a-googler committed 2 years ago

8a7bfe12

added Einsum layer to NNX

chiamp committed 2 years ago

792e390b

Merge pull request #3757 from carlosgmartin:contributing_git_fetch

a-googler committed 2 years ago

c2af6143

cgarciae committed 2 years ago

4160a136

[nnx] fix demo notebook

cgarciae committed 2 years ago

585e844d

Fix tests after applying JAX key-reuse checker. See:

chiamp committed 2 years ago

83d118ad

Older