google/flax

Pull Requests Commits

[nnx] refactor vmap

cgarciae committed 1 year ago

4b49ca85

Merge pull request #3960 from chiamp:fp8

a-googler committed 1 year ago

40c78bd7

chiamp committed 1 year ago

5e8e5951

Simplify extended dtypes rules part 1. Start by removing sharding specific rules from EDtypes. This is because we always want to replicate the trailing dims introduced by Edtypes.

yashk2810 committed 1 year ago

1619778d

Merge pull request #3906 from chiamp:nnx_transforms

a-googler committed 1 year ago

16a82100

Merge pull request #3923 from chiamp:haiku

a-googler committed 1 year ago

e37a37fb

Merge pull request #3912 from chiamp:rng

a-googler committed 1 year ago

6e29b41d

parameterized nnx_transforms tests

chiamp committed 1 year ago

0c6c7475

updated rng guide

chiamp committed 1 year ago

39584d7a

upgraded haiku guide to include NNX

chiamp committed 1 year ago

e926b179

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

IvyZX committed 1 year ago

3f47c1d9

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

a-googler committed 1 year ago

6e63497f

Add direct penzai.treescope support for NNX objects.

danieldjohnson committed 1 year ago

144486b5

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

a-googler committed 1 year ago

2cfd1747

[nnx] improve vmap axis size detection

cgarciae committed 1 year ago

90d656a8

[nnx] fix nnx_basics dependencies

cgarciae committed 1 year ago

e30754dd

Merge pull request #3941 from google:v0.8.5

a-googler committed 1 year ago

8c529b23

cgarciae committed 1 year ago

2ccb5555

Merge pull request #3936 from google:nnx-stabilize

a-googler committed 1 year ago

c84510d0

[nnx] move out of experimental

cgarciae committed 1 year ago

67fa0516

[nnx] fix UpdateContextManager

Cristian Garcia committed 1 year ago

b1cb9521

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

a-googler committed 1 year ago

9b90e98b

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

cgarciae committed 1 year ago

Verified 57dd8489

[nnx] improve update context

cgarciae committed 1 year ago

eeab2142

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

a-googler committed 1 year ago

312b2dd9

Merge pull request #3929 from IvyZX:lora

a-googler committed 1 year ago

b84be498

Merge pull request #3935 from chiamp:sequential

a-googler committed 1 year ago

45db1748

Merge pull request #3934 from chiamp:nnx_conv

a-googler committed 1 year ago

8bb8c871

[nnx] fix substate mutability

cgarciae committed 1 year ago

0741f192

Add LoRA to NNX

IvyZX committed 1 year ago

0f6dc0e6

Older