[JAX] Fix GEMM partitioning to reduce over nested contracting mesh axes - #3429

Merged

phu0ngng merged 4 commits into

NVIDIA:mainfrom

phu0ngng:jax/fix_gemm_partitioning

Aug 28, 2026

Commits on Aug 26, 2026

[JAX] Fix GEMM partitioning to reduce over nested contracting mesh axes
phu0ngng
committed
8ae2dfe
[JAX] Use CPU devices explicitly in gemm partitioning test mesh
phu0ngng
committed
bad3f05

Commits on Aug 27, 2026

[JAX] Expand GEMM partitioning tests with consistent FFN fwd/bwd sharding cases
phu0ngng
committed
f36703b

Commits on Aug 28, 2026

Merge branch 'main' into jax/fix_gemm_partitioning
phu0ngng
authored
dda83f8