google/flax

Pull Requests Commits

[nnx] update bridge

cgarciae committed 221 days ago

112ea170

Merge pull request #4820 from google:nnx-fix-rng-dtype-check

a-googler committed 226 days ago

8447c032

Merge pull request #4823 from google:fix-trailing-whitespaces

a-googler committed 226 days ago

4bea978c

Fixed .type deprecated property usage

vfdev-5 committed 226 days ago

539d9f01

Fix trailing whitespace

vfdev-5 committed 226 days ago

f179e8ee

Merge pull request #4800 from google:nnx-remove-variable-state

a-googler committed 226 days ago

2e779f57

Allow using LazyRngs for flax init/apply.

levskaya committed 226 days ago

1ccd1598

[nnx] fix Rngs dtype check

cgarciae committed 227 days ago

164c6be2

Merge pull request #4814 from divyashreepathihalli:use-id

a-googler committed 228 days ago

e147958f

divyashreepathihalli committed 228 days ago

a587aaca

Merge pull request #4813 from google:to-linen-eval-shape-2

a-googler committed 231 days ago

8c22f295

modify NNX to use id(variable) instead of nnx.Variables as dictionary keys

divyashreepathihalli committed 231 days ago

f411a322

Merge pull request #4776 from google:add-note-gemma-guide-about-model-license

a-googler committed 232 days ago

23d1a1b4

Merge pull request #4790 from google:enable-docs-nnx-remat

a-googler committed 232 days ago

a3f5be21

[nnx] ToLinen add abtract_init flag

cgarciae committed 232 days ago

7268e5f5

Merge pull request #4809 from IvyZX:doc2

a-googler committed 233 days ago

188baabc

Deduplicate contributing and philosophy and move to main site

IvyZX committed 233 days ago

103f65de

Preserve sharding information in axes_scan

oliverdutton-iso committed 233 days ago

4e2becd8

Internal change

kho committed 233 days ago

cbb4f80f

Merge pull request #4808 from google:to-linen-method-calls

a-googler committed 233 days ago

171b28e8

[nnx] allow method calls in ToLinen

cgarciae committed 234 days ago

61d0e9e0

[nnx] remove VariableState

cgarciae committed 237 days ago

e8984b70

Merge pull request #4755 from google:nnx-remove-all

a-googler committed 240 days ago

99f8e598

Merge pull request #4798 from google:ver-0.10.7

a-googler committed 240 days ago

1ecd34f9

update version to 0.10.7

cgarciae committed 240 days ago

d2cf8c2b

Fix typo in Flax `nnx_basics` doc.

a-googler committed 246 days ago

cca78723

Internal changes

basioli-k committed 246 days ago

6770f73f

Merge pull request #4775 from google:try-fix-colab-link-nnx

a-googler committed 247 days ago

7d22696c

Fixed nnx.remat docstring rendering

vfdev-5 committed 248 days ago

Verified 4870afe5

Merge pull request #4788 from google:add

a-googler committed 248 days ago

0b9bb4ee

Older