google/flax

Pull Requests Commits

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

a-googler committed 299 days ago

dbce607b

update ml-collections

cgarciae committed 299 days ago

ccdd29f7

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

IvyZX committed 301 days ago

aa897790

Merge pull request #4630 from google:optimize-jit

a-googler committed 302 days ago

1238a616

Merge pull request #4705 from IvyZX:einsum

a-googler committed 302 days ago

63106905

Add custom einsum op to Einsum()

IvyZX committed 303 days ago

000ee89b

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

Jan Hosang committed 303 days ago

32ad6c57

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

Jan Hosang committed 303 days ago

5bba7b45

[nnx] refactor GraphDef

cgarciae committed 303 days ago

8cb8a7a2

Add `attention_bias` parameter to `MultiHeadDotProductAttention`.

Jan Hosang committed 306 days ago

2fae1bad

Merge pull request #4627 from google:remove-deprecated

a-googler committed 307 days ago

aa9d5728

Merge pull request #4686 from kaixih:support_fp8_einsum

a-googler committed 308 days ago

fe0b35d7

kaixih committed 308 days ago

01cdcc74

kaixih committed 309 days ago

76f89b1f

Update the fp8 doc

kaixih committed 310 days ago

5242c7a0

Address comments

kaixih committed 310 days ago

1fd19374

Support fp8 einsum

kaixih committed 313 days ago

e0e253bd

Merge pull request #4681 from IvyZX:lora

a-googler committed 313 days ago

e58aeea9

Make LoRA's dtype arg useful

IvyZX committed 314 days ago

0ca02718

fix metadata sharing in nn.Variable.replace

liudangyi committed 314 days ago

e16cfc44

Automatic model selection for Gemma 3 models.

casaro committed 314 days ago

031d745a

Add support for layer-specific rope scale factors.

casaro committed 314 days ago

97ed31e4

Sow top activations based on absolute value.

casaro committed 314 days ago

5d41f6a4

Merge pull request #4656 from google:v0.10.5

a-googler committed 317 days ago

68b69f7f

cgarciae committed 317 days ago

01254afb

levskaya committed 320 days ago

6abddc9a

Merge pull request #4657 from IvyZX:wrapper-typo

a-googler committed 321 days ago

7600ad10

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

a-googler committed 321 days ago

d6ebb9df

ubuntu upgrade 24.04

IvyZX committed 322 days ago

Verified 983a0888

Fix lazy_init typo

IvyZX committed 322 days ago

Verified ce6e2496

Older