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

Add is_causal mask argument and tests

9bab1485

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 182 days ago

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

Reviewers

cgarciae

vfdev-5

Assignees

No one assigned

Labels

pull ready

Milestone

No milestone