google/flax

Pull Requests Commits

rajasekharporeddy committed 2 years ago

Verified 956ad6b3

rajasekharporeddy committed 2 years ago

Verified 89fd9f4d

Merge pull request #3993 from wenchenvincent:ocp_nanoo_fp8

a-googler committed 2 years ago

5b8265c6

Merge pull request #3808 from google:linen-allow-multi-compact

a-googler committed 2 years ago

646a6463

Merge pull request #3988 from google:nnx-lora-param

a-googler committed 2 years ago

b5503da5

[linen] allows multiple compact methods

cgarciae committed 2 years ago

1f291978

Minor refactor.

wenchenvincent committed 2 years ago

a6f52ae7

allow imagenet training profiling to be disabled in config

levskaya committed 2 years ago

afaa721e

Merge pull request #4040 from google:v0.8.6

a-googler committed 2 years ago

cb646c8b

Fixed test_fp8_meta_dtype.

wenchenvincent committed 2 years ago

bc083b83

Addressed the reviewer's comment. Avoided changing the existing

wenchenvincent committed 2 years ago

541bddf4

Added support of NANOO fp8.

wenchenvincent committed 2 years ago

28bdb790

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

a-googler committed 2 years ago

15e0e8da

[nnx] add Using Filters guide

cgarciae committed 2 years ago

0aa5fe6f

cgarciae committed 2 years ago

2698262e

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

[nnx] LoRAParam inherits from Param

cgarciae committed 2 years ago

24c44f6c

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

Newer Older