flax
Added is_causal mask argument to flax.nnx.dot_product_attention
#5093

Merged

Added is_causal mask argument to flax.nnx.dot_product_attention #5093

copybara-service merged 1 commit into google:main from nnx-sdpa-parity

vfdev-5 commented on 2025-11-17

ibbyml force pushed from b06c6fc7 to 0c334257 55 days ago

Add is_causal mask argument and tests

9bab1485

ibbyml force pushed from 0c334257 to 9bab1485 55 days ago

vfdev-5 approved these changes on 2025-11-21

vfdev-5 commented on 2025-12-01

cgarciae approved these changes on 2025-12-10

cgarciae added pull ready

copybara-service merged be1db789 into main 33 days ago

ibbyml deleted the nnx-sdpa-parity branch 32 days ago

Reviewers

cgarciae

vfdev-5

Assignees

No one assigned

Labels

pull ready

Milestone

No milestone