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

[phu0ngng]

Description

Please include a brief summary of the changes, relevant motivation and context.

Fixes # (issue)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• Change A
• Change B

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Greptile Summary

The PR updates JAX GEMM partition inference to reduce across every mesh axis shared by both operands’ contracting dimensions while gathering non-shared axes. It also adds CPU-based partitioning tests for Megatron TP, sequence-parallel TP, and compatibility cases.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/jax/cpp_extensions/gemm.py	Extends contracting-axis handling from scalar specifications to nested tuples and retains only axes participating in the reduction.
tests/jax/test_gemm_partitioning.py	Adds CPU mesh tests for nested contracting-axis reductions and fixes the previously reported default-backend mesh dependency by selecting CPU devices explicitly.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  L[LHS contracting axes] --> I[Shared-axis intersection]
  R[RHS contracting axes] --> I
  I --> REDUCE[Reduction specification]
  L --> FILTER[Retain shared axes]
  R --> FILTER
  FILTER --> GEMM[Partitioned GEMM]
  GEMM --> REDUCE

Loading

_{Reviews (3): Last reviewed commit: "Merge branch 'main' into jax/fix_gemm_pa..." | Re-trigger Greptile}

[xiaopoc]

Tested this PR on DeepSeek-V3 with 32 H100 GPUs using TP8, DP4, EP16, FSDP2, and ETP1. The previous global WGrad buffer/all-gather disappeared, HLO now shows WGrad reduce-scatter over the tensor axis, and the training loss decreases normally. This fixes the issue I observed, thanks.

[xiaopoc]

Looks good to me. Verified in the test described in my comment above.

[jberchtold-nvidia]

Overall looks good, thanks! Left a couple small comments and a question

[phu0ngng]

/te-ci L1 JAX

[jberchtold-nvidia]

LGTM, thanks!