google/flax

Pull Requests Commits

Port gemma/layers to NNX

sibonaci committed 2 years ago

f62da76f

Helper function for loading params from a linen module

sibonaci committed 2 years ago

ab8c773a

Merge pull request #4006 from IvyZX:notftest

a-googler committed 2 years ago

8a2e6413

remove tf from jax_utils_test.py to fix head pytest

IvyZX committed 2 years ago

44e48f9c

Cristian Garcia committed 2 years ago

fbdc3394

Merge pull request #3991 from chiamp:iter

a-googler committed 2 years ago

e6b83adc

added nnx.graph.iter_children

chiamp committed 2 years ago

098a2d52

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

Older