google/flax

Pull Requests Commits

Merge pull request #4180 from google:nnx-fix-optimizer

a-googler committed 2 years ago

671130b3

Merge pull request #4171 from IvyZX:bdg-guide

a-googler committed 2 years ago

d373856d

[nnx] improve Optimizer metadata propagation

cgarciae committed 2 years ago

4c472878

Add bridge guide

IvyZX committed 2 years ago

cdfa0a73

Move all NNX content up a level to be equal with Linen, to make python packaging more consistent.

IvyZX committed 2 years ago

71b964b4

Fixes a small bug in flax.linen.share_scope, where the scopes of children of the module being merged that were created before setup(),were not being updated to point to the new scope, and so they would end up staying under the original tree.

marksandler2 committed 2 years ago

05353e64

Merge pull request #4168 from google:nnx-improve-messaging

a-googler committed 2 years ago

b494c3f6

Internal change

IvyZX committed 2 years ago

e848a995

[nnx] disallow Array leaves

Cristian Garcia committed 2 years ago

90715bef

[nnx] improve landing page and nnx_basics messaging

cgarciae committed 2 years ago

0be4e143

Merge pull request #4170 from google:nnx-variable-proxy

a-googler committed 2 years ago

aded9ac4

Merge pull request #4135 from 8bitmp3:update-nnx-surgery

a-googler committed 2 years ago

bc594fbc

Upgrade Flax NNX Model Surgery

8bitmp3 committed 2 years ago

4c36c8ae

Merge pull request #3922 from wenscarl:direct_quant

a-googler committed 2 years ago

c44b9169

[nnx] add Variable proxy methods

cgarciae committed 2 years ago

6d7ad155

Remove unnecessary warnings.

wenscarl committed 2 years ago

e9d19aca

Add type: ignore

wenscarl committed 2 years ago

b4dc0949

wenscarl committed 2 years ago

4f5a722e

Direct quantization for FP8 Dense Layer.

wenscarl committed 2 years ago

6f5cee12

Merge pull request #4157 from google:nnx-improve-scan

a-googler committed 2 years ago

35c4edf8

[nnx] improve nnx.scan in_axes/out_axes

cgarciae committed 2 years ago

aeeda794

Merge pull request #4155 from google:nnx-fix-nnx-basics-2

a-googler committed 2 years ago

98dff5e8

[nnx] use jax-like transforms API in nnx_basics

cgarciae committed 2 years ago

1a78acb5

Merge pull request #4158 from IvyZX:bdg-doc

a-googler committed 2 years ago

559ba024

Add nnx bridge API reference to site

IvyZX committed 2 years ago

93821f7d

Merge pull request #4145 from IvyZX:bdg

a-googler committed 2 years ago

839db8c1

Support linen <-> nnx metadata box converging in bridge

IvyZX committed 2 years ago

4bb2152d

Merge pull request #4149 from google:nnx-fix-nnx-basics-prints

a-googler committed 2 years ago

afbc5021

[nnx] clear nnx basics pip logs

cgarciae committed 2 years ago

366aeff3

update Flax to version 0.9.0

Cristian Garcia committed 2 years ago

fc2e5993

Older