[JAX] Refactor and trim TE JAX Attn testing#2542

Merged

KshitijLakhani merged 7 commits intoNVIDIA:mainfrom

KshitijLakhani:klakhani/maint/reduce-test-time

Jan 9, 2026

Commits on Dec 24, 2025

Pick a leaner set of combinations for TE JAX CP attn tests such that only those cp,dp,tp combinations are picked where cp*dp*tp is equal to num gpus
KshitijLakhani
committed