google/flax

Pull Requests Commits

remove VariableState

cgarciae committed 255 days ago

0c3838ae

[nnx] make Variable a pytree

cgarciae committed 263 days ago

2f3b344f

Fixes a bug in type annotations for scope.param (unbox=True should accept callable[..., T | AxisMEtadata[T]] and return T, while unbox=False should always return the same thing as what callable returning.

marksandler2 committed 263 days ago

54fad0a8

Merge pull request #4652 from jlperla:identity

a-googler committed 264 days ago

240a5107

Merge pull request #4724 from google:0.10.6

a-googler committed 264 days ago

0fde6365

Merge pull request #4639 from google:flaxlib-types

a-googler committed 264 days ago

dc062922

[nnx] add more flaxlib types

cgarciae committed 264 days ago

334f3830

Lower-bounded using new JAX release

jlperla committed 265 days ago

5a4ec98b

Merge branch 'google:main' into identity

jlperla committed 265 days ago

Verified f3c11518

remove "nnx" collection in ToLinen

Cristian Garcia committed 265 days ago

d3ec212f

cgarciae committed 265 days ago

40d7df36

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

a-googler committed 269 days ago

dbce607b

update ml-collections

cgarciae committed 269 days ago

ccdd29f7

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

IvyZX committed 271 days ago

aa897790

Merge pull request #4630 from google:optimize-jit

a-googler committed 272 days ago

1238a616

Merge pull request #4705 from IvyZX:einsum

a-googler committed 272 days ago

63106905

Add custom einsum op to Einsum()

IvyZX committed 272 days ago

000ee89b

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

Jan Hosang committed 273 days ago

32ad6c57

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

Jan Hosang committed 273 days ago

5bba7b45

[nnx] refactor GraphDef

cgarciae committed 273 days ago

8cb8a7a2

Add `attention_bias` parameter to `MultiHeadDotProductAttention`.

Jan Hosang committed 276 days ago

2fae1bad

Merge pull request #4627 from google:remove-deprecated

a-googler committed 277 days ago

aa9d5728

Merge pull request #4686 from kaixih:support_fp8_einsum

a-googler committed 278 days ago

fe0b35d7

kaixih committed 278 days ago

01cdcc74

kaixih committed 279 days ago

76f89b1f

Update the fp8 doc

kaixih committed 279 days ago

5242c7a0

Address comments

kaixih committed 280 days ago

1fd19374

Support fp8 einsum

kaixih committed 283 days ago

e0e253bd

Merge pull request #4681 from IvyZX:lora

a-googler committed 283 days ago

e58aeea9

Make LoRA's dtype arg useful

IvyZX committed 283 days ago

0ca02718

Older