google/flax

Pull Requests Commits

[nnx] add jit donation test

cgarciae committed 2 years ago

c0236ea6

Merge pull request #4038 from google:python-3.10

a-googler committed 2 years ago

f610e6ce

upgrade python to 3.10 + use pyupgrade

cgarciae committed 2 years ago

97ca05d2

Port gemma/sampler to NNX

sibonaci committed 2 years ago

fed77561

Port gemma/transformer to NNX

sibonaci committed 2 years ago

c50d1533

Merge pull request #4005 from IvyZX:surgery

a-googler committed 2 years ago

0c19b6b7

Make `force_fp32_for_softmax` arg in `MultiHeadDotProductAttention` useful.

IvyZX committed 2 years ago

9e6b8193

Merge pull request #3996 from jakevdp:oldest-jax

a-googler committed 2 years ago

3b218708

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

levskaya committed 2 years ago

e5cb2f76

CI: add test run against oldest supported jax version

jakevdp committed 2 years ago

58020c41

Merge pull request #4002 from sanderland:fixpickle

a-googler committed 2 years ago

791df52d

Add NNX surgery guide

IvyZX committed 2 years ago

7eeaa7ba

Merge pull request #4025 from IvyZX:unbreak

a-googler committed 2 years ago

9b6748aa

Fix Gemma test to unbreak head

IvyZX committed 2 years ago

1c4e3489

Unpin the Orbax pip version.

IvyZX committed 2 years ago

68d5c816

Internal change to fix current head CI

IvyZX committed 2 years ago

bda7da1e

Port gemma/modules to NNX

sibonaci committed 2 years ago

9ca8db56

Support Python 3.9

sibonaci committed 2 years ago

1edfcfa5

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

a-googler committed 2 years ago

4d9081e5

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

a-googler committed 2 years ago

54fe15f2

cgarciae committed 2 years ago

1e0a5d10

Port gemma/layers to NNX

sibonaci committed 2 years ago

f62da76f

Helper function for loading params from a linen module

sibonaci committed 2 years ago

ab8c773a

[nnx] add PathContains Filter

cgarciae committed 2 years ago

5ee8e98b

Sander Land committed 2 years ago

f89718e2

Merge pull request #4006 from IvyZX:notftest

a-googler committed 2 years ago

8a2e6413

remove tf from jax_utils_test.py to fix head pytest

IvyZX committed 2 years ago

44e48f9c

Sander Land committed 2 years ago

255564ed

edit flax/errors.py

Sander Land committed 2 years ago

95bad7a1

Sander Land committed 2 years ago

d9731cc3

Older