google/flax

Pull Requests Commits

[linen] dynamic rngs

cgarciae committed 1 year ago

9fa88f93

Update PyO3 to 0.21.2

a-googler committed 1 year ago

fec32840

Merge pull request #4183 from google:nnx-fsdp-toy-example

a-googler committed 1 year ago

2ad97319

[nnx] add nnx.variables + FSDP toy example

cgarciae committed 1 year ago

f7d3873b

Merge pull request #4255 from 8bitmp3:add-graph-flatten-autosummary

a-googler committed 1 year ago

9f52f662

Merge pull request #4267 from google:remove-flax-nnx-urls

a-googler committed 1 year ago

2a2187bd

remove flax-nnx urls

cgarciae committed 1 year ago

b12c4b41

Merge pull request #4266 from google:custom-vjp-docs

a-googler committed 1 year ago

601a9150

Merge pull request #4226 from google:nnx-train-state-toy-example

a-googler committed 1 year ago

62a3414f

[nnx] add custom_vjp to docs

cgarciae committed 1 year ago

312368e6

Merge pull request #4256 from Vilin97:main

a-googler committed 1 year ago

40b9b3a8

Merge pull request #4253 from google:cloud-pickle-support

a-googler committed 1 year ago

b487b52e

Merge pull request #4240 from google:why-nnx

a-googler committed 1 year ago

94f6440c

add Why NNX doc

cgarciae committed 1 year ago

caa061d0

Merge pull request #4264 from 8bitmp3:fixes-nnx-2

a-googler committed 1 year ago

11f36163

Add checkpointing guide to website index.

IvyZX committed 1 year ago

bb4d44af

Light update to Flax NNX Transforms doc

8bitmp3 committed 1 year ago

a1543a81

Merge pull request #4249 from IvyZX:ckpt-guide

a-googler committed 1 year ago

9a3a1fc1

Merge pull request #4251 from ariG23498:patch-1

a-googler committed 1 year ago

551942e9

Fix typo: `impost` to `import`

Vilin97 committed 1 year ago

a792633d

Update graph.rst

8bitmp3 committed 1 year ago

Verified 149144e0

Add flatten to nnx.graph autosummary in graph.rst

8bitmp3 committed 1 year ago

Verified 34b8c6b0

[nnx] add cloudpickle support

cgarciae committed 1 year ago

c55cd448

Merge pull request #4250 from google:enable-notebook-doctest

a-googler committed 1 year ago

5d314528

Update index.rst

ariG23498 committed 1 year ago

Verified fcaa303f

enable doctest on notebooks

cgarciae committed 1 year ago

146d2ff3

Merge pull request #4248 from 8bitmp3:fix-links

a-googler committed 1 year ago

2d64500c

Merge pull request #4244 from IvyZX:nnx-landing

a-googler committed 1 year ago

a8bca5d2

add checkpointing guide

IvyZX committed 1 year ago

742c9262

Nitting and adding links

8bitmp3 committed 1 year ago

ac3fafba

Older