google/flax

Pull Requests Commits

fix site banners

cgarciae committed 2 years ago

4d11e98e

Partially revert #4192 which sets back a bunch of previous merged pushes.

IvyZX committed 2 years ago

9eb0a619

Merge pull request #4192 from google:nnx-transforms-guide

a-googler committed 2 years ago

03e034d0

cgarciae committed 2 years ago

ea57438f

Merge pull request #4191 from google:nnx-optimize-jit

a-googler committed 2 years ago

d111adf7

cgarciae committed 2 years ago

2079f15b

Merge pull request #4186 from IvyZX:errors

a-googler committed 2 years ago

b9679643

Merge pull request #4189 from IvyZX:axis-cleanup

a-googler committed 2 years ago

61ea8a66

clean up axis hooks

IvyZX committed 2 years ago

978bc3dd

Merge nnx.errors to flax.errors

IvyZX committed 2 years ago

e8b0a1fd

Merge pull request #4161 from IvyZX:bdg-logic

a-googler committed 2 years ago

f948154d

Support linen.LogicallyPartitioned <-> nnx.Variable

IvyZX committed 2 years ago

a9cb80b8

Merge pull request #4185 from google:nnx-transform-metadata-issue

a-googler committed 2 years ago

1b72435a

[nnx] enable sharding transformation on integer prefixes

cgarciae committed 2 years ago

93c06270

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

Older