google/flax

Pull Requests Commits

[nnx] improving performance of `_graph_flatten` and `_graph_unflatten` functions

erfanzar committed 1 year ago

bf36dec6

Merge pull request #4400 from google:nnx-more-optimizations-p2

a-googler committed 1 year ago

5d896bc1

Merge pull request #4402 from tilakrayal:patch-2

a-googler committed 1 year ago

be2f8615

[nnx] optimize Variable

cgarciae committed 1 year ago

7b50ffec

Update haiku_to_flax.rst

tilakrayal committed 1 year ago

Verified 18b59dde

Merge pull request #4399 from google:nnx-more-optimizations

a-googler committed 1 year ago

fe8fb59e

[nnx] optimize NodeDef.attributes

cgarciae committed 1 year ago

861936f7

Automated Code Change: Suppress pytype errors

Martin Huschenbett committed 1 year ago

6bc9858a

Merge pull request #4384 from google:nnx-performance-guide

a-googler committed 1 year ago

abc11552

[nnx] add performance guide notebook

cgarciae committed 1 year ago

94793f9b

Merge pull request #4386 from 8bitmp3:update-flax-readme

a-googler committed 1 year ago

be26138c

Upgrade Flax readme to NNX

8bitmp3 committed 1 year ago

bddf5f24

Merge pull request #4390 from IvyZX:loop

a-googler committed 1 year ago

4e258985

Fix fori_loop and while_loop on multiple modules

IvyZX committed 1 year ago

19a4abf0

Merge pull request #4387 from google:v0.10.2

a-googler committed 1 year ago

d89c955d

cgarciae committed 1 year ago

52a8c482

Merge pull request #4377 from google:nnx-pytree-optimization

a-googler committed 1 year ago

6382a2bf

[nnx] add some optimizations to graph.py

cgarciae committed 1 year ago

ba575246

[nnx] use HashableMapping instead of FrozenDict

cgarciae committed 1 year ago

7fd97dc4

Merge pull request #4370 from 8bitmp3:update-nnx-fori_loop-docstring

a-googler committed 1 year ago

91bf758d

Merge pull request #4371 from 8bitmp3:update-nnx-while_loop

a-googler committed 1 year ago

9147a7cd

Merge pull request #4381 from google:nnx-checkify

a-googler committed 1 year ago

d13f047e

[nnx] add checkify

cgarciae committed 1 year ago

ec73cfcf

Merge pull request #4373 from 8bitmp3:add-nnx-remat-docstring

a-googler committed 1 year ago

856b6b58

Add flax.nnx.remat docstring

8bitmp3 committed 1 year ago

8fbe99cf

Setup the flaxlib in C++, using Meson and Nanobind.

IvyZX committed 1 year ago

72f4971c

Merge pull request #4374 from 8bitmp3:add-nnx-eval_shap-docstring

a-googler committed 1 year ago

f265a5e2

Merge pull request #4379 from google:nnx-fix-fori-loop

a-googler committed 1 year ago

ac3e85a9

[nnx] fix while_loop/fori_loop bug when sharing references

cgarciae committed 1 year ago

c48905a2

Merge pull request #4376 from google:nnx-benchmark

a-googler committed 1 year ago

480a1960

Older