flax
Use jax.nn.dot_product_attention when possible
#4649

Merged

Use jax.nn.dot_product_attention when possible #4649

copybara-service merged 1 commit into google:main from IvyZX:attn

IvyZX requested a review from

cgarciae 274 days ago

Use jax.nn.dot_product_attention when possible

6f7322db

IvyZX force pushed from c521915a to 6f7322db 274 days ago

cgarciae approved these changes on 2025-03-25

cgarciae added pull ready

copybara-service merged 527cb4a4 into main 273 days ago

Reviewers

cgarciae

Assignees

No one assigned

Labels

pull ready

Milestone

No milestone