google/flax

Pull Requests Commits

Adding Counter metric.

Alex Alemi committed 2 years ago

20fe6132

Merge pull request #3975 from chiamp:training

a-googler committed 2 years ago

998792a4

Merge pull request #3980 from chiamp:docstrings

a-googler committed 2 years ago

0e5f920e

updated nnx.training

chiamp committed 2 years ago

31adb009

Merge pull request #3981 from chiamp:np_bool

a-googler committed 2 years ago

5b7d671f

Merge pull request #3982 from jkarwowski:jkarwowski-patch-1

a-googler committed 2 years ago

5d6ca397

Fixes for batch norm docs

jkarwowski committed 2 years ago

Verified 8b871506

updated rnglib docstring

chiamp committed 2 years ago

d8d4e5d0

fix deprecation warning

chiamp committed 2 years ago

884a1ca3

* Adds `pmap` and `Pmap`. `static_broadcasted_argnums`, `donate_argnums`, and `global_arg_shapes` are not yet supported.

Cristian Garcia committed 2 years ago

78d85af6

Merge pull request #3958 from chiamp:docstrings

a-googler committed 2 years ago

3f56060e

updated nnx.graph docstrings

chiamp committed 2 years ago

fb986f65

Merge pull request #3974 from chiamp:rnglib_doc

a-googler committed 2 years ago

3ef8cd84

updated nnx.Conv and nnx.ConvTranspose

chiamp committed 2 years ago

d5e7b3dd

Merge pull request #3966 from chiamp:module_doc

a-googler committed 2 years ago

7c6a6555

make flax work with upcoming JAX change to tree_map (being more careful about

mattjj committed 2 years ago

e9bde53c

Merge pull request #3972 from chiamp:nnx_nn

a-googler committed 2 years ago

860bef59

Merge pull request #3964 from google:nnx-generalize-state

a-googler committed 2 years ago

89c08028

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

Older