google/flax

Pull Requests Commits

Internal change

IvyZX committed 2 years ago

adb2c8dd

[nnx] disallow Array leaves

Cristian Garcia committed 2 years ago

90715bef

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

Merge pull request #4148 from google:fix-headers

a-googler committed 2 years ago

1b08425c

remove repeated license headers

cgarciae committed 2 years ago

4be609d7

Merge pull request #4141 from google:nnx-landing-page

a-googler committed 2 years ago

a0622b07

[nnx] improve landing page

cgarciae committed 2 years ago

4e17aa10

Merge pull request #3963 from google:nnx-experimental-transforms

a-googler committed 2 years ago

3a9d8339

Merge pull request #4123 from cool-RR:patch-4

a-googler committed 2 years ago

d0942f9b

Merge pull request #4137 from IvyZX:wrapper-fix

a-googler committed 2 years ago

1b2cca9f

Merge pull request #4136 from google:use-uv

a-googler committed 2 years ago

2ebd3984

Older