google/flax

Pull Requests Commits

Internal change

IvyZX committed 260 days ago

8ee48cd1

Merge pull request #4686 from kaixih:support_fp8_einsum

a-googler committed 260 days ago

fe0b35d7

kaixih committed 260 days ago

01cdcc74

kaixih committed 261 days ago

76f89b1f

Update the fp8 doc

kaixih committed 262 days ago

5242c7a0

Address comments

kaixih committed 262 days ago

1fd19374

Support fp8 einsum

kaixih committed 265 days ago

e0e253bd

Merge pull request #4681 from IvyZX:lora

a-googler committed 265 days ago

e58aeea9

Make LoRA's dtype arg useful

IvyZX committed 266 days ago

0ca02718

fix metadata sharing in nn.Variable.replace

liudangyi committed 266 days ago

e16cfc44

Automatic model selection for Gemma 3 models.

casaro committed 267 days ago

031d745a

Add support for layer-specific rope scale factors.

casaro committed 267 days ago

97ed31e4

Sow top activations based on absolute value.

casaro committed 267 days ago

5d41f6a4

Merge pull request #4656 from google:v0.10.5

a-googler committed 269 days ago

68b69f7f

cgarciae committed 269 days ago

01254afb

levskaya committed 273 days ago

6abddc9a

Merge pull request #4657 from IvyZX:wrapper-typo

a-googler committed 273 days ago

7600ad10

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

a-googler committed 273 days ago

d6ebb9df

ubuntu upgrade 24.04

IvyZX committed 274 days ago

Verified 983a0888

Fix lazy_init typo

IvyZX committed 274 days ago

Verified ce6e2496

Fix flaky vmap test tolerance.

levskaya committed 275 days ago

0dec8a80

Merge pull request #4649 from IvyZX:attn

a-googler committed 275 days ago

527cb4a4

Use jax.nn.dot_product_attention when possible

IvyZX committed 276 days ago

6f7322db

Merge pull request #4648 from IvyZX:fix-head

a-googler committed 276 days ago

5f7af869

roll up package jax version and uv.lock

IvyZX committed 276 days ago

6a77c8a9

Merge pull request #4612 from google:nnx-array-leaves

a-googler committed 276 days ago

097fffce

Merge pull request #4646 from google/release-webhook

ZacCranko committed 276 days ago

Verified ba59c336

[nnx] support Array leaves in graph nodes

cgarciae committed 276 days ago

dfd1788c

discord release webhook

cgarciae committed 276 days ago

9e4d3e3f

Stop passing reduce_axes to jax.grad, jax.vjp, and jax.value_and_grad.

jburnim committed 279 days ago

f18baf07

Older