google/flax

Pull Requests Commits

fix restore_checkpoint example in docstring

Cristian Garcia committed 1 year ago

46dc51eb

Merge pull request #4079 from danielwatson6:scope-shape-check

a-googler committed 1 year ago

d718655f

Merge pull request #4084 from google:nnx-split-merge-inputs

a-googler committed 1 year ago

8a880812

Merge pull request #4099 from google:nnx-reseed

a-googler committed 1 year ago

a7bdadba

Merge pull request #4081 from IvyZX:bridge

a-googler committed 1 year ago

d8bc1947

Add shape-based lazy init to LinenToNNX (prev LinenWrapper)

IvyZX committed 1 year ago

0139c90e

[nnx] add reseed

cgarciae committed 1 year ago

3f8c62fb

[nnx] add split/merge_inputs

cgarciae committed 1 year ago

82993122

[linen] fold rngs on jit to improve caching

Cristian Garcia committed 1 year ago

afcaf66d

[linen] generalize transform caching

Cristian Garcia committed 1 year ago

4e83e091

Merge pull request #4098 from google:fix-main

a-googler committed 1 year ago

65cd19e7

fix mypy failures due tu numpy update

cgarciae committed 1 year ago

30b0b0a2

Merge pull request #4092 from IvyZX:unbreak

a-googler committed 1 year ago

ceacc090

Ignore Orbax warning in deprecated flax.training and unbreak head

IvyZX committed 1 year ago

b0ebefcf

Merge pull request #4004 from google:nnx-pure

a-googler committed 1 year ago

c4066cce

Merge pull request #4059 from google:nnx-improve-filters-guide

a-googler committed 1 year ago

62127f81

[nnx] improve filters guide

cgarciae committed 1 year ago

7937ca08

Merge pull request #4077 from knightXun:patch-1

a-googler committed 1 year ago

93babfd2

Merge pull request #4078 from google:nnx-extract-apis

a-googler committed 1 year ago

e7f43dd7

[nnx] add extract APIs

cgarciae committed 1 year ago

16c2c324

Merge pull request #4076 from kaixih:add_fp8_doc

a-googler committed 1 year ago

e8021e9f

Perform shape checks for self.param AFTER unboxing, make shmap users happy

danielwatson6 committed 1 year ago

bc689087

[example]: remove lm1b useless parallism rules

knightXun committed 1 year ago

Verified 93291277

Add doc for fp8

kaixih committed 1 year ago

350d9f83

Reverting unsafe_pytree #4030

Cristian Garcia committed 1 year ago

239d4e6b

cgarciae committed 1 year ago

f89e62cf

Remove some outdated backward-compatibility code.

IvyZX committed 1 year ago

5c971436

Modify Flax checkpointing in preparation for cl/650338576.

cpgaffney1 committed 1 year ago

e59b14e9

Merge pull request #4070 from google:nnx-improve-mnist-tutorial

a-googler committed 1 year ago

1b58348d

Merge pull request #4066 from IvyZX:bridge

a-googler committed 1 year ago

359ec7fd

Older