google/flax

Pull Requests Commits

add fast_flatten

cgarciae committed 357 days ago

53896f53

[nnx] add more flaxlib types

cgarciae committed 357 days ago

39dca1e7

Merge pull request #4716 from google:update-ml-collections

a-googler committed 361 days ago

dbce607b

update ml-collections

cgarciae committed 361 days ago

ccdd29f7

Make fully replicated array before saving checkpoints for examples that use pmap.

IvyZX committed 363 days ago

aa897790

Merge pull request #4630 from google:optimize-jit

a-googler committed 364 days ago

1238a616

Merge pull request #4705 from IvyZX:einsum

a-googler committed 364 days ago

63106905

Add custom einsum op to Einsum()

IvyZX committed 1 year ago

000ee89b

Rollback of attention_bias parameter, because the change overrides the attention bias for injected attention functions.

Jan Hosang committed 1 year ago

32ad6c57

Unit tests for `attention_bias` parameter to `MultiHeadDotProductAttention`. Add parameter to all overloads to make pytype happy.

Jan Hosang committed 1 year ago

5bba7b45

[nnx] refactor GraphDef

cgarciae committed 1 year ago

8cb8a7a2

Add `attention_bias` parameter to `MultiHeadDotProductAttention`.

Jan Hosang committed 1 year ago

2fae1bad

Merge pull request #4627 from google:remove-deprecated

a-googler committed 1 year ago

aa9d5728

Merge pull request #4686 from kaixih:support_fp8_einsum

a-googler committed 1 year ago

fe0b35d7

kaixih committed 1 year ago

01cdcc74

kaixih committed 1 year ago

76f89b1f

Update the fp8 doc

kaixih committed 1 year ago

5242c7a0

Address comments

kaixih committed 1 year ago

1fd19374

Support fp8 einsum

kaixih committed 1 year ago

e0e253bd

Merge pull request #4681 from IvyZX:lora

a-googler committed 1 year ago

e58aeea9

Make LoRA's dtype arg useful

IvyZX committed 1 year ago

0ca02718

fix metadata sharing in nn.Variable.replace

liudangyi committed 1 year ago

e16cfc44

Automatic model selection for Gemma 3 models.

casaro committed 1 year ago

031d745a

Add support for layer-specific rope scale factors.

casaro committed 1 year ago

97ed31e4

Sow top activations based on absolute value.

casaro committed 1 year ago

5d41f6a4

Merge pull request #4656 from google:v0.10.5

a-googler committed 1 year ago

68b69f7f

cgarciae committed 1 year ago

01254afb

levskaya committed 1 year ago

6abddc9a

Merge pull request #4657 from IvyZX:wrapper-typo

a-googler committed 1 year ago

7600ad10

Merge pull request #4659 from IvyZX:IvyZX-patch-1

a-googler committed 1 year ago

d6ebb9df

Older