"""TP primitive ops (ADR-0027 D6).

``copy_to_tp_region`` / ``reduce_from_tp_region`` are forward-only in the
initial scope (backward pass is future work). ``scatter`` / ``gather`` are
not implemented — they require an all-gather kernel that is not yet
available in KernBench (see ADR-0027 D9).
"""
from __future__ import annotations

from typing import Any


def copy_to_tp_region(x: Any) -> Any:
    """Forward: identity. Backward: all-reduce. (Training is future.)"""
    return x


def reduce_from_tp_region(x: Any, torch: Any) -> Any:
    """Forward: all-reduce. Backward: identity."""
    torch.distributed.all_reduce(x, op="sum")
    return x


def scatter_to_tp_region(x: Any) -> Any:
    raise NotImplementedError(
        "scatter_to_tp_region deferred — caller should create the sharded "
        "tensor directly (ADR-0027 D9)"
    )


def gather_from_tp_region(x: Any) -> Any:
    raise NotImplementedError(
        "gather_from_tp_region deferred — requires all-gather kernel (ADR-0027 D9)"
    )
