google/flax

Pull Requests Commits

Merge pull request #4851 from google:update-support-mutable-array

a-googler committed 201 days ago

065d6e6b

[nnx] add update support for mutable array

cgarciae committed 201 days ago

d1c1043f

Merge pull request #4842 from google:optimizer-keep-model

a-googler committed 201 days ago

aeed717a

Merge pull request #4842 from google:optimizer-keep-model

a-googler committed 201 days ago

94c94802

Fix re-export and docs for identity

jlperla committed 201 days ago

5518913d

create ModelAndOptimizer symbol

Cristian Garcia committed 201 days ago

77196487

[nnx] add Optimizer keep_model param

cgarciae committed 202 days ago

c6498d57

Merge pull request #4811 from google:nnx-pytree-modules-2

a-googler committed 202 days ago

ae208ee7

Add split_rngs docstring

vfdev-5 committed 202 days ago

8f6ab14e

Migrate `transforms_test.py` from `.value` to `[...]`

lukeyeh committed 204 days ago

726a01dd

[nnx] add flax_pytree_module flag

cgarciae committed 204 days ago

50e043c2

Merge pull request #4822 from google:add-train-script-gemma-example

a-googler committed 204 days ago

97f4a499

Merge pull request #4815 from lukeyeh:push-tpqklpwqzxkq

a-googler committed 204 days ago

db0e3020

Added training script for Gemma model

vfdev-5 committed 210 days ago

2eb7baa1

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

a-googler committed 210 days ago

8447c032

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

a-googler committed 210 days ago

4bea978c

Fixed .type deprecated property usage

vfdev-5 committed 210 days ago

539d9f01

Fix trailing whitespace

vfdev-5 committed 210 days ago

f179e8ee

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

a-googler committed 211 days ago

2e779f57

Allow using LazyRngs for flax init/apply.

levskaya committed 211 days ago

1ccd1598

[nnx] fix Rngs dtype check

cgarciae committed 211 days ago

164c6be2

Merge pull request #4814 from divyashreepathihalli:use-id

a-googler committed 212 days ago

e147958f

refactor: move usages of `.value` to `[...]` in modules_test.py

lukeyeh committed 212 days ago

9acdcfdc

divyashreepathihalli committed 212 days ago

a587aaca

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

a-googler committed 215 days ago

8c22f295

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

divyashreepathihalli committed 215 days ago

f411a322

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

a-googler committed 216 days ago

23d1a1b4

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

a-googler committed 216 days ago

a3f5be21

[nnx] ToLinen add abtract_init flag

cgarciae committed 216 days ago

7268e5f5

Merge pull request #4809 from IvyZX:doc2

a-googler committed 217 days ago

188baabc

Newer Older