[JAX] Fix counting of synced and wrapped (async, fusion) collectives in HLO

[aybchan]

Description

Follow-up to #3398

With collective fusion, XLA may now emit an all-reduce behind an asynchronous fusion wrapper in HLO:

%fusion-start = ((f32[]), f32[]) async-start(...),
  calls=%async_computation

The parsing logic previously recognized fusion-start as an async start instruction, but inferred the collective type from its name. Because fusion-start contains neither all-reduce nor all-gather, the communication bytes were incorrectly classified as other.

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

• [JAX] Fix counting of synced and wrapped (async, fusion) collectives in HLO tests

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 the JAX distributed-test HLO parser to recognize ROOT-prefixed and synchronous collectives and to resolve collectives hidden behind async or fusion wrappers.

• Strips the optional ROOT prefix before instruction classification.
• Resolves a wrapper's calls= computation and classifies embedded all-reduce or all-gather operations.
• Uses opcode-aware matching for synchronous collectives.

Confidence Score: 5/5

The PR appears safe to merge because no concrete changed-code failure remains supported by the available repository evidence.

The updated parser preserves direct async classification while adding handling for ROOT syntax, synchronous opcodes, and called computations used by fusion wrappers.

Important Files Changed

Filename	Overview
tests/jax/distributed_test_base.py	Extends collective classification for synchronous, ROOT-prefixed, and wrapped HLO instructions without an established actionable defect.

_{Reviews (1): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile}

[jberchtold-nvidia]

/te-ci L1 jax

[jberchtold-nvidia]

Change LGTM pending CI. I will merge when CI completes. Thanks Alex!