google/flax

Pull Requests Commits

[nnx] keras integration

cgarciae committed 1 year ago

c79476f0

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

[nnx] fix ToLinen kwargs

Cristian Garcia committed 1 year ago

86ff7af3

[nnx] use HashableMapping instead of FrozenDict

cgarciae committed 1 year ago

347211bc

Add flax.nnx.eval_shape docstring

8bitmp3 committed 1 year ago

Verified ddd7847b

Add flax.nnx.eval_shape docstring

8bitmp3 committed 1 year ago

Verified 6b5b3006

Lint flax.nnx.while_loop docstring

8bitmp3 committed 1 year ago

Verified 5495f4e4

Lint flax.nnx.fori_loop docstring

8bitmp3 committed 1 year ago

Verified 84fa22ea

Merge pull request #4350 from IvyZX:logical-axis

a-googler committed 1 year ago

d31f2906

Merge pull request #4354 from google:nnx-data-parallel-example

a-googler committed 1 year ago

803aca8b

[nnx] add data parallel toy example

cgarciae committed 1 year ago

a2237caf

Merge pull request #4352 from google:nnx-pure-dict-support

a-googler committed 1 year ago

0fabd3fd

Avoid assert_array_equal for PRNG keys.

Jake VanderPlas committed 1 year ago

83eebe27

[nnx] support pure dicts

cgarciae committed 1 year ago

97c3e9c1

support passing arguments directly to the struct.dataclass decorator

Qwlouse committed 1 year ago

c9cebeee

Fix PRNG handling in `nn.jit` under `nn.scan`.

danielsuo committed 1 year ago

0f631a27

Merge pull request #4200 from 8bitmp3:upgrade-haiku-linen-to-nnx

a-googler committed 1 year ago

0679702d

Older