google/flax

Pull Requests Commits
[nnx] fix grad
cgarciae committed 1 year ago
1e0a5d10
Port gemma/layers to NNX
sibonaci committed 1 year ago
f62da76f
Helper function for loading params from a linen module
sibonaci committed 1 year ago
ab8c773a
[nnx] add PathContains Filter
cgarciae committed 1 year ago
5ee8e98b
rename
Sander Land committed 1 year ago
f89718e2
Merge pull request #4006 from IvyZX:notftest
a-googler committed 1 year ago
8a2e6413
remove tf from jax_utils_test.py to fix head pytest
IvyZX committed 1 year ago
44e48f9c
test
Sander Land committed 1 year ago
255564ed
edit flax/errors.py
Sander Land committed 1 year ago
95bad7a1
fix pickle
Sander Land committed 1 year ago
d9731cc3
[nnx] fix vmap
Cristian Garcia committed 1 year ago
fbdc3394
Merge pull request #3991 from chiamp:iter
a-googler committed 1 year ago
e6b83adc
added nnx.graph.iter_children
chiamp committed 1 year ago
098a2d52
Merge pull request #3983 from chiamp:optimizer
a-googler committed 1 year ago
56216fa6
added wrt option to nnx.Optimizer.update
chiamp committed 1 year ago
9a40b3fa
Merge pull request #3989 from google:nnx-fix-vmap
a-googler committed 1 year ago
551eaf02
Merge pull request #3986 from chiamp:variables
a-googler committed 1 year ago
7a69dcd1
[nnx] vectorize vmap split counts
cgarciae committed 1 year ago
e4f417ac
updated nnx.variables docstrings
chiamp committed 1 year ago
7ca2ee14
Merge pull request #3975 from chiamp:training
a-googler committed 1 year ago
998792a4
Merge pull request #3980 from chiamp:docstrings
a-googler committed 1 year ago
0e5f920e
updated nnx.training
chiamp committed 1 year ago
31adb009
Merge pull request #3981 from chiamp:np_bool
a-googler committed 1 year ago
5b7d671f
Merge pull request #3982 from jkarwowski:jkarwowski-patch-1
a-googler committed 1 year ago
5d6ca397
Fixes for batch norm docs
jkarwowski committed 1 year ago
Verified 8b871506
updated rnglib docstring
chiamp committed 1 year ago
d8d4e5d0
fix deprecation warning
chiamp committed 1 year ago
884a1ca3
* Adds `pmap` and `Pmap`. `static_broadcasted_argnums`, `donate_argnums`, and `global_arg_shapes` are not yet supported.
Cristian Garcia committed 1 year ago
78d85af6
Merge pull request #3958 from chiamp:docstrings
a-googler committed 1 year ago
3f56060e
updated nnx.graph docstrings
chiamp committed 1 year ago
fb986f65
Terms Of Use Privacy Policy Refund Policy Impressum