google/flax

Pull Requests Commits

Merge pull request #4186 from IvyZX:errors

a-googler committed 1 year ago

b9679643

Merge pull request #4189 from IvyZX:axis-cleanup

a-googler committed 1 year ago

61ea8a66

clean up axis hooks

IvyZX committed 1 year ago

978bc3dd

Merge nnx.errors to flax.errors

IvyZX committed 1 year ago

e8b0a1fd

Merge pull request #4161 from IvyZX:bdg-logic

a-googler committed 1 year ago

f948154d

Support linen.LogicallyPartitioned <-> nnx.Variable

IvyZX committed 1 year ago

a9cb80b8

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

a-googler committed 1 year ago

1b72435a

[nnx] enable sharding transformation on integer prefixes

cgarciae committed 1 year ago

93c06270

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

a-googler committed 1 year ago

671130b3

Merge pull request #4171 from IvyZX:bdg-guide

a-googler committed 1 year ago

d373856d

[nnx] improve Optimizer metadata propagation

cgarciae committed 1 year ago

4c472878

Add bridge guide

IvyZX committed 1 year ago

cdfa0a73

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

IvyZX committed 1 year 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 1 year ago

05353e64

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

a-googler committed 1 year ago

b494c3f6

Internal change

IvyZX committed 1 year ago

e848a995

[nnx] disallow Array leaves

Cristian Garcia committed 1 year ago

90715bef

[nnx] improve landing page and nnx_basics messaging

cgarciae committed 1 year ago

0be4e143

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

a-googler committed 1 year ago

aded9ac4

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

a-googler committed 1 year ago

bc594fbc

Upgrade Flax NNX Model Surgery

8bitmp3 committed 1 year ago

4c36c8ae

Merge pull request #3922 from wenscarl:direct_quant

a-googler committed 1 year ago

c44b9169

[nnx] add Variable proxy methods

cgarciae committed 1 year ago

6d7ad155

Remove unnecessary warnings.

wenscarl committed 1 year ago

e9d19aca

Add type: ignore

wenscarl committed 1 year ago

b4dc0949

wenscarl committed 1 year ago

4f5a722e

Direct quantization for FP8 Dense Layer.

wenscarl committed 1 year ago

6f5cee12

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

a-googler committed 1 year ago

35c4edf8

[nnx] improve nnx.scan in_axes/out_axes

cgarciae committed 1 year ago

aeeda794

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

a-googler committed 1 year ago

98dff5e8

Older