google/flax

Pull Requests Commits

[nnx] fix pre-commit

cgarciae committed 2 years ago

abf36e5b

Minor grammar fixes.

mcsmart76 committed 2 years ago

Verified bc9a3165

Rename all the NNX tests to internal naming & build conventions.

IvyZX committed 2 years ago

3f47c1d9

Merge pull request #3942 from google:fix-nnx-basics-notebook

a-googler committed 2 years ago

6e63497f

Add direct penzai.treescope support for NNX objects.

danieldjohnson committed 2 years ago

144486b5

Merge pull request #3947 from google:nnx-vmap-improve-axes-size-detection

a-googler committed 2 years ago

2cfd1747

[nnx] improve vmap axis size detection

cgarciae committed 2 years ago

90d656a8

[nnx] fix nnx_basics dependencies

cgarciae committed 2 years ago

e30754dd

Merge pull request #3941 from google:v0.8.5

a-googler committed 2 years ago

8c529b23

cgarciae committed 2 years ago

2ccb5555

Merge pull request #3936 from google:nnx-stabilize

a-googler committed 2 years ago

c84510d0

[nnx] move out of experimental

cgarciae committed 2 years ago

67fa0516

[nnx] fix UpdateContextManager

Cristian Garcia committed 2 years ago

b1cb9521

Merge pull request #3933 from google:nnx-improve-update-context

a-googler committed 2 years ago

9b90e98b

Update flax/experimental/nnx/nnx/graph.py

cgarciae committed 2 years ago

Verified 57dd8489

[nnx] improve update context

cgarciae committed 2 years ago

eeab2142

Merge pull request #3932 from google:nnx-fix-state-getitem

a-googler committed 2 years ago

312b2dd9

Merge pull request #3929 from IvyZX:lora

a-googler committed 2 years ago

b84be498

Merge pull request #3935 from chiamp:sequential

a-googler committed 2 years ago

45db1748

Merge pull request #3934 from chiamp:nnx_conv

a-googler committed 2 years ago

8bb8c871

[nnx] fix substate mutability

cgarciae committed 2 years ago

0741f192

Add LoRA to NNX

IvyZX committed 2 years ago

0f6dc0e6

added test for nnx.ConvTranspose

chiamp committed 2 years ago

a4c0d437

added equivalence test for

chiamp committed 2 years ago

93a7c79a

Merge pull request #3927 from google:nnx-cond

a-googler committed 2 years ago

7542b28a

[nnx] transforms refactor

cgarciae committed 2 years ago

6b31a980

Merge pull request #3921 from google:nnx-compact

a-googler committed 2 years ago

2bdcee76

[nnx] add compat

cgarciae committed 2 years ago

72ea9599

Merge pull request #3919 from google:nnx-iternodes-docs

a-googler committed 2 years ago

f2b2ea33

[nnx] add iter_graph

cgarciae committed 2 years ago

ed49ab36

Older