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

Draft

KshitijLakhani wants to merge 2 commits intoNVIDIA:mainfrom

KshitijLakhani:klakhani/maint/reduce-test-time

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
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot]
committed