google/flax

Pull Requests Commits

Port gemma sampling tutorial

sibonaci committed 1 year ago

dfc9ad09

Merge pull request #4038 from google:python-3.10

a-googler committed 1 year ago

f610e6ce

upgrade python to 3.10 + use pyupgrade

cgarciae committed 1 year ago

97ca05d2

Port gemma/sampler to NNX

sibonaci committed 1 year ago

fed77561

Port gemma/transformer to NNX

sibonaci committed 1 year ago

c50d1533

Merge pull request #4005 from IvyZX:surgery

a-googler committed 1 year ago

0c19b6b7

Make `force_fp32_for_softmax` arg in `MultiHeadDotProductAttention` useful.

IvyZX committed 1 year ago

9e6b8193

Merge pull request #3996 from jakevdp:oldest-jax

a-googler committed 1 year ago

3b218708

Call user-defined variable transforms before determining axis size in nn.vmap.

levskaya committed 1 year ago

e5cb2f76

CI: add test run against oldest supported jax version

jakevdp committed 1 year ago

58020c41

Merge pull request #4002 from sanderland:fixpickle

a-googler committed 1 year ago

791df52d

Add NNX surgery guide

IvyZX committed 1 year ago

7eeaa7ba

Merge pull request #4025 from IvyZX:unbreak

a-googler committed 1 year ago

9b6748aa

Fix Gemma test to unbreak head

IvyZX committed 1 year ago

1c4e3489

Unpin the Orbax pip version.

IvyZX committed 1 year ago

68d5c816

Internal change to fix current head CI

IvyZX committed 1 year ago

bda7da1e

Port gemma/modules to NNX

sibonaci committed 1 year ago

9ca8db56

Support Python 3.9

sibonaci committed 1 year ago

1edfcfa5

Merge pull request #4011 from google:nnx-path-filter

a-googler committed 1 year ago

4d9081e5

Merge pull request #4007 from google:nnx-fix-grad

a-googler committed 1 year ago

54fe15f2

cgarciae committed 1 year ago

1e0a5d10

Port gemma/layers to NNX

sibonaci committed 1 year ago

f62da76f

Helper function for loading params from a linen module

sibonaci committed 1 year ago

ab8c773a

[nnx] add PathContains Filter

cgarciae committed 1 year ago

5ee8e98b

Sander Land committed 1 year ago

f89718e2

Merge pull request #4006 from IvyZX:notftest

a-googler committed 1 year ago

8a2e6413

remove tf from jax_utils_test.py to fix head pytest

IvyZX committed 1 year ago

44e48f9c

Sander Land committed 1 year ago

255564ed

edit flax/errors.py

Sander Land committed 1 year ago

95bad7a1

Sander Land committed 1 year ago

d9731cc3

Older