google/flax

Pull Requests Commits

[linen] allows multiple compact methods

cgarciae committed 1 year ago

1f291978

Minor refactor.

wenchenvincent committed 1 year ago

a6f52ae7

allow imagenet training profiling to be disabled in config

levskaya committed 1 year ago

afaa721e

Merge pull request #4040 from google:v0.8.6

a-googler committed 1 year ago

cb646c8b

Fixed test_fp8_meta_dtype.

wenchenvincent committed 1 year ago

bc083b83

Addressed the reviewer's comment. Avoided changing the existing

wenchenvincent committed 1 year ago

541bddf4

Added support of NANOO fp8.

wenchenvincent committed 1 year ago

28bdb790

Merge pull request #4028 from google:nnx-filters-guide2

a-googler committed 1 year ago

15e0e8da

[nnx] add Using Filters guide

cgarciae committed 1 year ago

0aa5fe6f

cgarciae committed 1 year ago

2698262e

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

[nnx] LoRAParam inherits from Param

cgarciae committed 1 year ago

24c44f6c

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

Newer Older