google/flax

Pull Requests Commits

[nnx] add more vmap tests

cgarciae committed 2 years ago

ed87eb0e

Merge pull request #3983 from chiamp:optimizer

a-googler committed 2 years ago

56216fa6

added wrt option to nnx.Optimizer.update

chiamp committed 2 years ago

9a40b3fa

Merge pull request #3989 from google:nnx-fix-vmap

a-googler committed 2 years ago

551eaf02

Merge pull request #3986 from chiamp:variables

a-googler committed 2 years ago

7a69dcd1

[nnx] vectorize vmap split counts

cgarciae committed 2 years ago

e4f417ac

updated nnx.variables docstrings

chiamp committed 2 years ago

7ca2ee14

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

Older