google/flax

Pull Requests Commits

Merge pull request #3634 from chiamp:batch_apply

a-googler committed 1 year ago

77da09e9

Remove calls to deprecated function jax.lax.tie_in

Jake VanderPlas committed 1 year ago

993aa2c5

Copybara import of the project:

chiamp committed 1 year ago

0370a610

Copybara import of the project:

a-googler committed 1 year ago

4d075e53

Merge pull request #3621 from google:nontransparent

a-googler committed 2 years ago

7e9bb761

added BatchApply

chiamp committed 2 years ago

6454502a

add compact_name_scope

cgarciae committed 2 years ago

c9e63a7b

support transform reuse for lift.jit decorator

Cristian Garcia committed 2 years ago

93d3eb78

Merge pull request #3625 from chiamp:ruff

a-googler committed 2 years ago

c6c169a3

disable ruff linter

chiamp committed 2 years ago

37605208

Merge pull request #3590 from jianyizh:jianyi/embed

a-googler committed 2 years ago

3824622c

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

a-googler committed 2 years ago

a32cf1d2

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

a-googler committed 2 years 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

Older