google/flax

Pull Requests Commits

[nnx] add unsafe_graph_cache

cgarciae committed 2 years ago

cfdf1d41

Merge pull request #4353 from IvyZX:conds

a-googler committed 2 years ago

591cd402

IvyZX committed 2 years ago

6ad09a3d

Merge pull request #4345 from google:v0.10.1

a-googler committed 2 years ago

f740ab3c

Merge pull request #4343 from IvyZX:conds

a-googler committed 2 years ago

72341579

Add conditionals `nnx.while_loop` and `nnx.switch`

IvyZX committed 2 years ago

4f5d6fb9

Merge pull request #4346 from google:update-state-docstrings

a-googler committed 2 years ago

8292d9cc

Merge pull request #4349 from 8bitmp3:define-model-surgery

a-googler committed 2 years ago

b8bdafbb

Define Flax NNX Model Surgery in docs

8bitmp3 committed 2 years ago

817755f3

Add support for storing arbitrary PyTrees with `Module.perturb()`

duckworthd committed 2 years ago

917c0974

[nnx] update State and variables docstrings

cgarciae committed 2 years ago

13b40770

cgarciae committed 2 years ago

300dcc60

Merge pull request #4306 from google:nnx-custom-vjp-issue

a-googler committed 2 years ago

e4dad9ca

[nnx] fix custom_vjp

cgarciae committed 2 years ago

78034830

Remove non-lazy RNG compat mode and flag from flax.

levskaya committed 2 years ago

b2d23e48

Merge pull request #4334 from google:cleanup-gemma

a-googler committed 2 years ago

bcc12b5d

Merge pull request #4333 from google:improve-transforms-guide

a-googler committed 2 years ago

58d1e537

[nnx] cleanup gemma notebook

cgarciae committed 2 years ago

158bcea2

[nnx] improve transforms guide

cgarciae committed 2 years ago

a8e2ab1f

Merge pull request #4338 from IvyZX:docfix

a-googler committed 2 years ago

ab122af7

Add API reference for flax.nnx.nn and improve landing page

IvyZX committed 2 years ago

e1415041

Merge pull request #4337 from 8bitmp3:fixup-haiku-to-flax

a-googler committed 2 years ago

0f044c8b

Fix typos in Flax NNX Migrating from Haiku to Flax

8bitmp3 committed 2 years ago

b83dce71

Merge pull request #4332 from jakevdp:fix-lift-test

a-googler committed 2 years ago

cd5db467

Remove remaining references to Penzai from NNX.

danieldjohnson committed 2 years ago

3895d51b

Avoid assert_array_equal on PRNG keys.

jakevdp committed 2 years ago

27c06b62

Avoid using float32 in normalization for mean/var and scale/bias parameters when force_float32_reductions=False

eunjaekim-0 committed 2 years ago

40f080ed

Remove GeGLU activation function and golden tests.

levskaya committed 2 years ago

ad90642b

Merge pull request #4331 from levskaya:fix_lint

a-googler committed 2 years ago

85ca2c0e

Merge pull request #4324 from jakevdp:rm-compat

a-googler committed 2 years ago

1e885003

Older