google/flax

Pull Requests Commits
Merge pull request #4351 from jlperla:lbfgs_support
a-googler committed 1 year ago
d8b1a921
Add logical axis global context support for NNX
IvyZX committed 1 year ago
4d3006af
Merge pull request #4353 from IvyZX:conds
a-googler committed 1 year ago
591cd402
add
IvyZX committed 1 year ago
6ad09a3d
Merge pull request #4345 from google:v0.10.1
a-googler committed 1 year ago
f740ab3c
Merge pull request #4343 from IvyZX:conds
a-googler committed 1 year ago
72341579
Add conditionals `nnx.while_loop` and `nnx.switch`
IvyZX committed 1 year ago
4f5d6fb9
Moved everything to kwargs
jlperla committed 1 year ago
342adde2
Merge pull request #4346 from google:update-state-docstrings
a-googler committed 1 year ago
8292d9cc
Punctuation
jlperla committed 1 year ago
cea33dac
Renamed tests
jlperla committed 1 year ago
0daea7d0
Add LBFGS arguments
jlperla committed 1 year ago
aef15427
Merge pull request #4349 from 8bitmp3:define-model-surgery
a-googler committed 1 year ago
b8bdafbb
Define Flax NNX Model Surgery in docs
8bitmp3 committed 1 year ago
817755f3
Add support for storing arbitrary PyTrees with `Module.perturb()`
duckworthd committed 1 year ago
917c0974
[nnx] update State and variables docstrings
cgarciae committed 1 year ago
13b40770
v0.10.1
cgarciae committed 1 year ago
300dcc60
Merge pull request #4306 from google:nnx-custom-vjp-issue
a-googler committed 1 year ago
e4dad9ca
[nnx] fix custom_vjp
cgarciae committed 1 year ago
78034830
Remove non-lazy RNG compat mode and flag from flax.
levskaya committed 1 year ago
b2d23e48
Merge pull request #4334 from google:cleanup-gemma
a-googler committed 1 year ago
bcc12b5d
Merge pull request #4333 from google:improve-transforms-guide
a-googler committed 1 year ago
58d1e537
[nnx] cleanup gemma notebook
cgarciae committed 1 year ago
158bcea2
[nnx] improve transforms guide
cgarciae committed 1 year ago
a8e2ab1f
Merge branch 'main' into upgrade-haiku-linen-to-nnx
8bitmp3 committed 1 year ago
Verified 84fa3167
Upgrade Flax NNX Migrating from Haiku doc
8bitmp3 committed 1 year ago
4ed2ba8c
Merge pull request #4338 from IvyZX:docfix
a-googler committed 1 year ago
ab122af7
Add API reference for flax.nnx.nn and improve landing page
IvyZX committed 1 year ago
e1415041
Merge pull request #4337 from 8bitmp3:fixup-haiku-to-flax
a-googler committed 1 year ago
0f044c8b
Fix typos in Flax NNX Migrating from Haiku to Flax
8bitmp3 committed 1 year ago
b83dce71
Terms Of Use Privacy Policy Refund Policy Impressum