"""Cross-source correlation modules.

This module provides functionality for:
- Cross-rank clock alignment using collective sync points
- Matching collectives across ranks by type, size, ordering
- Linking PyTorch operators to NCCL/RCCL collectives
"""

from wafer_core.lib.distributed_traces.correlation.collective_matcher import (
    CollectiveMatcher,
    match_collectives_across_ranks,
)
from wafer_core.lib.distributed_traces.correlation.pytorch_nccl_linker import (
    PyTorchNCCLLinker,
    link_pytorch_to_nccl,
)
from wafer_core.lib.distributed_traces.correlation.rank_aligner import (
    RankAligner,
    align_ranks,
)

__all__ = [
    # Rank alignment
    "RankAligner",
    "align_ranks",
    # Collective matching
    "CollectiveMatcher",
    "match_collectives_across_ranks",
    # PyTorch-NCCL linking
    "PyTorchNCCLLinker",
    "link_pytorch_to_nccl",
]
