"""Eager peer-copy oracle: PyTorch copies/events, no raw peer pointers or spin protocol.

Debug only, outside graph capture. FP32 addition then optional BF16 rounding matches the push collective.
"""
from __future__ import annotations

import torch


def all_reduce(parts: tuple[torch.Tensor, torch.Tensor], dtype: torch.dtype = torch.bfloat16
               ) -> tuple[torch.Tensor, torch.Tensor]:
    if any(x.dtype != torch.float32 or not x.is_cuda for x in parts) or parts[0].shape != parts[1].shape:
        raise ValueError("peer-copy reference needs two same-shaped FP32 CUDA tensors")
    ready = []
    for x in parts:
        with torch.cuda.device(x.device):
            if torch.cuda.is_current_stream_capturing():
                raise RuntimeError("peer-copy reference is eager only")
            event = torch.cuda.Event()
            event.record(torch.cuda.current_stream(x.device))
            ready.append(event)
    result = []
    for r, x in enumerate(parts):
        with torch.cuda.device(x.device):
            torch.cuda.current_stream(x.device).wait_event(ready[1 - r])
            peer = parts[1 - r].to(x.device, non_blocking=True)
            result.append((x + peer).to(dtype))
    return result[0], result[1]