google/flax

Pull Requests Commits

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

Merge pull request #3733 from google:nnx-remove-flags-context

a-googler committed 2 years ago

0280160c

Add git fetch upstream to contributing doc.

carlosgmartin committed 2 years ago

84809ef9

Merge pull request #3712 from carlosgmartin:conv_lstm_cell_docs

a-googler committed 2 years ago

3da44bc1

Add ConvLSTMCell to docs.

carlosgmartin committed 2 years ago

1c9cdc1f

Merge pull request #3752 from chiamp:intercept

a-googler committed 2 years ago

50b589f4

[nnx] remove flagslib

cgarciae committed 2 years ago

2206c40a

Merge pull request #3732 from google:add-pytree-support

a-googler committed 2 years ago

9fbf69dc

fixed intercept_methods docstring

chiamp committed 2 years ago

fca6b454

removed getattr/setattr magic from nnx.Pytree

chiamp committed 2 years ago

c279f248

[nnx] add pytree support

cgarciae committed 2 years ago

c8987746

Merge pull request #3753 from IvyZX:orbax-quickfix

a-googler committed 2 years ago

ce8a3c74

Fix doctest and unbreak head

IvyZX committed 2 years ago

3b23dc8c

Fix instances of JAX key reuse in flax tests

Jake VanderPlas committed 2 years ago

6f2b08e0

Merge pull request #3728 from google:nnx-cleanup-graph-utils

a-googler committed 2 years ago

0c7d256d

Merge pull request #3742 from chiamp:nnx_dataclass

a-googler committed 2 years ago

e8e0b2df

removed nnx dataclass

chiamp committed 2 years ago

d4ac6c3f

Merge pull request #3623 from kaixih:fp8_meta_custom_grad_accumulation

a-googler committed 2 years ago

e4282ee1

Newer Older