google/flax

Pull Requests Commits

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

Merge pull request #4351 from jlperla:lbfgs_support

a-googler committed 1 year ago

d8b1a921

Add logical axis global context support for NNX

IvyZX committed 1 year ago

4d3006af

Merge pull request #4353 from IvyZX:conds

a-googler committed 1 year ago

591cd402

IvyZX committed 1 year ago

6ad09a3d

Merge pull request #4345 from google:v0.10.1

a-googler committed 1 year ago

f740ab3c

Merge pull request #4343 from IvyZX:conds

a-googler committed 1 year ago

72341579

Add conditionals `nnx.while_loop` and `nnx.switch`

IvyZX committed 1 year ago

4f5d6fb9

Moved everything to kwargs

jlperla committed 1 year ago

342adde2

Merge pull request #4346 from google:update-state-docstrings

a-googler committed 1 year ago

8292d9cc

jlperla committed 1 year ago

cea33dac

jlperla committed 1 year ago

0daea7d0

Add LBFGS arguments

jlperla committed 1 year ago

aef15427

Merge pull request #4349 from 8bitmp3:define-model-surgery

a-googler committed 1 year ago

b8bdafbb

Define Flax NNX Model Surgery in docs

8bitmp3 committed 1 year ago

817755f3

Add support for storing arbitrary PyTrees with `Module.perturb()`

duckworthd committed 1 year ago

917c0974

[nnx] update State and variables docstrings

cgarciae committed 1 year ago

13b40770

cgarciae committed 1 year ago

300dcc60

Merge pull request #4306 from google:nnx-custom-vjp-issue

a-googler committed 1 year ago

e4dad9ca

[nnx] fix custom_vjp

cgarciae committed 1 year ago

78034830

Remove non-lazy RNG compat mode and flag from flax.

levskaya committed 1 year ago

b2d23e48

Merge pull request #4334 from google:cleanup-gemma

a-googler committed 1 year ago

bcc12b5d

Newer Older