google/flax

Pull Requests Commits

support transform reuse for lift.jit decorator

Cristian Garcia committed 1 year ago

93d3eb78

Merge pull request #3625 from chiamp:ruff

a-googler committed 1 year ago

c6c169a3

disable ruff linter

chiamp committed 1 year ago

37605208

Merge pull request #3590 from jianyizh:jianyi/embed

a-googler committed 1 year ago

3824622c

Merge pull request #3620 from google:fix-lm1b-2

a-googler committed 1 year ago

a32cf1d2

Merge pull request #3609 from google:nnx-flaglib-get

a-googler committed 1 year ago

8cce2704

jianyizh committed 2 years ago

9aee6602

Merge pull request #3619 from google:expose-module-functions-2

a-googler committed 2 years ago

d58e6dde

fix lm1b data sharding

cgarciae committed 2 years ago

360a0c2f

[#3572](https://github.com/google/flax/pull/3572) was rolled back because of internal test breakages. This PR re-adds the attention changes again and fixes internal tests

chiamp committed 2 years ago

9cfe078a

add missing docs for module functions

cgarciae committed 2 years ago

202ba4e2

Rollback of Copybara import of the project:

oskopek committed 2 years ago

ed44f521

Merge pull request #3572 from chiamp:attention_subclass

a-googler committed 2 years ago

90e6b606

Merge pull request #3610 from chiamp:quickstart

a-googler committed 2 years ago

2db554b5

fixed broken links on quick start guide

chiamp committed 2 years ago

5856a45b

added multiheadattention alias

chiamp committed 2 years ago

09f6ea47

[nnx] use positional-only args in Rngs and Flags

cgarciae committed 2 years ago

ba777c33

Merge pull request #3607 from google:nnx-improve-first-from

a-googler committed 2 years ago

fae3a30b

Merge pull request #3608 from google:nnx-improve-callable-proxy

a-googler committed 2 years ago

f63808cd

[nnx] cleanup CallableProxy

cgarciae committed 2 years ago

44e96a22

[nnx] improve runtime flags

cgarciae committed 2 years ago

55466275

Merge pull request #3576 from chiamp:nnx_dynamic_slice

a-googler committed 2 years ago

7cbc3c13

fixed NNX decode and dynamic slicing

chiamp committed 2 years ago

f17b4fe3

Merge pull request #3600 from google:nnx-clean-flaglib

a-googler committed 2 years ago

cb6843f2

Merge pull request #3581 from google:nnx-submodule-iterator

a-googler committed 2 years ago

3b23351e

delete flaglib duplicated copyright comment

cgarciae committed 2 years ago

039845ce

[nnx] add submodule iterator

cgarciae committed 2 years ago

1e75509a

Merge pull request #3591 from shoyer:patch-1

a-googler committed 2 years ago

85eb8c05

Merge pull request #3582 from google:nnx-flagslib-add-get-overloads

a-googler committed 2 years ago

cff22d3f

Merge pull request #3583 from google:nnx-add-missing-symbols

a-googler committed 2 years ago

bfe04e56

Older