google/flax

Pull Requests Commits

Modify Welford metric to return mean value.

Alex Alemi committed 2 years ago

e7be2872

[nnx] make State generic

cgarciae committed 2 years ago

bfb7efb5

updated NNX attention and nn docstrings

chiamp committed 2 years ago

4a10da73

Replace `XLACompatibleSharding` with `jax.sharding.Sharding` since the former is deprecated and will be removed in the future.

yashk2810 committed 2 years ago

21231af9

Adding Welford metric.

Alex Alemi committed 2 years ago

617a6cd7

updated nnx.module docstrings

chiamp committed 2 years ago

74534daa

Merge pull request #3928 from NeilGirdhar:use_mapping

a-googler committed 2 years ago

bff079b7

Merge pull request #3953 from mcsmart76:patch-1

a-googler committed 2 years ago

e5a19160

[nnx] fix pre-commit

cgarciae committed 2 years ago

abf36e5b

Merge pull request #3960 from chiamp:fp8

a-googler committed 2 years ago

40c78bd7

chiamp committed 2 years 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 2 years ago

1619778d

Merge pull request #3906 from chiamp:nnx_transforms

a-googler committed 2 years ago

16a82100

Merge pull request #3923 from chiamp:haiku

a-googler committed 2 years ago

e37a37fb

Merge pull request #3912 from chiamp:rng

a-googler committed 2 years ago

6e29b41d

parameterized nnx_transforms tests

chiamp committed 2 years ago

0c6c7475

updated rng guide

chiamp committed 2 years ago

39584d7a

upgraded haiku guide to include NNX

chiamp committed 2 years ago

e926b179

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

Make FlatState a Mapping instead of a dict

NeilGirdhar committed 2 years ago

ecb7cf37

[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

Newer Older