Diff Coverage

# limitations under the License.
# ============================================================================
"""SiTU eager reference and first-order autograd over CANN ops-nn."""

import importlib
from collections.abc import Callable
from typing import Any, Optional

# Torch is the native framework; the legacy platform-agnostic lint rule is obsolete here.
# pylint: disable=forbidden-backend-import
import torch
from torch.autograd.function import once_differentiable


def _situ_eager(gate: torch.Tensor, up: torch.Tensor, beta: float,
                linear_beta: Optional[float]) -> torch.Tensor:
    gate_float, up_float = gate.float(), up.float()
    activated_gate = beta * torch.tanh(gate_float / beta) * torch.sigmoid(gate_float)
    if linear_beta is not None and linear_beta > 0:
        up_float = linear_beta * torch.tanh(up_float / linear_beta)
    return (activated_gate * up_float).to(gate.dtype)


def _load_situ_ops() -> tuple[Callable, Callable]:
    """Resolve the optional ops-nn training APIs without hiding native errors."""
    try:
        # CANN is optional for CPU and eager-only callers.
        ops = importlib.import_module("cann_ops_nn")
    except (ImportError, OSError) as error:
        raise ImportError(
            "SiTU backend='ascendc' could not load cann_ops_nn. Install the standalone ops-nn Python package "
            "with compatible PyTorch, torch-npu and CANN libraries. "
            f"See hyper_parallel/components/functional/README.md. Original error: {error}"
        ) from error
    situ_glu = getattr(ops, "situ_glu", None)
    situ_glu_grad = getattr(ops, "situ_glu_grad", None)
    if not callable(situ_glu) or not callable(situ_glu_grad):
        raise ImportError(
            "SiTU backend='ascendc' requires cann_ops_nn.situ_glu and cann_ops_nn.situ_glu_grad. "
            "Install a package containing both training APIs and build both native operators. "
            "See hyper_parallel/components/functional/README.md."
        )
    return situ_glu, situ_glu_grad


class _SituAutograd(torch.autograd.Function):
    """Connect ops-nn's separate SituGlu and SituGluGrad entry points."""

    @staticmethod
    def forward(ctx: Any, packed: torch.Tensor, beta: float, linear_beta: float) -> torch.Tensor:
        """Run dense gate-first rows and retain the native backward inputs.

        Args:
            ctx: Autograd context holding inputs for backward.

        Returns:
            Activated tensor with half the input's last dimension.
        """
        situ_glu, situ_glu_grad = _load_situ_ops()
        rows = packed.reshape(-1, packed.shape[-1]).contiguous()
        ctx.save_for_backward(rows)
        ctx.input_shape = packed.shape
        ctx.beta = beta
        ctx.linear_beta = linear_beta
        ctx.situ_glu_grad = situ_glu_grad
        output = situ_glu(rows, dim=-1, beta=beta, linear_beta=linear_beta, activate_left=True)
        return output.reshape(*packed.shape[:-1], packed.shape[-1] // 2)

    @staticmethod
    @once_differentiable
    def backward(ctx: Any, grad_output: torch.Tensor) -> tuple[torch.Tensor, None, None]:
        """Restore the packed gradient with the same saturation parameters.

        Args:
            ctx: Autograd context populated by forward.

        Returns:
            Packed input gradient and no gradients for the scalar attributes.
        """
        (rows,) = ctx.saved_tensors
        grad_rows = grad_output.reshape(-1, rows.shape[-1] // 2).contiguous()
        gradient = ctx.situ_glu_grad(
            grad_rows, rows, dim=-1, beta=ctx.beta, linear_beta=ctx.linear_beta, activate_left=True,
        )
        return gradient.reshape(ctx.input_shape), None, None


def _situ_fused(packed: torch.Tensor, beta: float, linear_beta: float) -> torch.Tensor:
    return _SituAutograd.apply(packed, beta, linear_beta)


def situ(
    gate: torch.Tensor,
    up: torch.Tensor,
    beta: float = 1.0,
    linear_beta: Optional[float] = None,

    Returns:
        Activated values with gate's shape and dtype.
    """
    # Packing would lose the original split if gate and up had unequal widths.
    if gate.shape != up.shape:
        raise ValueError("SiTU gate and up must have identical shapes")
    if backend == "eager":
        return _situ_eager(gate, up, beta, linear_beta)
    return situ_packed(torch.cat((gate, up), dim=-1), beta, linear_beta, backend)


def situ_packed(
    packed: torch.Tensor,
    beta: float = 1.0,
    linear_beta: Optional[float] = None,
    backend: str = "eager",

    Returns:
        Tensor of shape [..., H] in the input dtype. Empty inputs keep connected
        gradients through tensor operations without launching a native kernel.
    """
    if backend not in ("eager", "ascendc"):
        raise ValueError("SiTU backend must be 'eager' or 'ascendc'")
    linear_beta = 0.0 if linear_beta is None else linear_beta
    if backend == "eager" or packed.numel() == 0:
        gate, up = packed.unflatten(-1, (2, packed.shape[-1] // 2)).unbind(-2)
        return _situ_eager(gate, up, beta, linear_beta)
    return _situ_fused(packed, beta, linear_beta)

# limitations under the License.
# ============================================================================
"""Block attention residuals with eager math and the ops-transformer fused API."""

import importlib
from collections.abc import Callable

# Torch is the native framework; the legacy platform-agnostic lint rule is obsolete here.
# pylint: disable=forbidden-backend-import
import torch


def _block_attention_residuals_eager(prefix: torch.Tensor, history: torch.Tensor,
                                    projection: torch.Tensor, norm_weight: torch.Tensor,
                                    norm_eps: float) -> torch.Tensor:
    values = torch.cat((history, prefix.unsqueeze(1)), dim=1)
    values_fp32 = values.float()
    normalized = values_fp32 * torch.rsqrt(values_fp32.square().mean(-1, keepdim=True) + norm_eps)
    query = projection.float().reshape(-1) * norm_weight.float()
    probabilities = (normalized * query).sum(-1).softmax(-1)
    return (probabilities.unsqueeze(-1) * values_fp32).sum(1).to(values.dtype)


def _load_block_attention_residuals() -> Callable:
    """Resolve the differentiable ops-transformer API and its backward dependency."""
    try:
        # CANN is optional for CPU and eager-only callers.
        ops = importlib.import_module("cann_ops_transformer")
    except (ImportError, OSError) as error:
        raise ImportError(
            "AttnRes backend='ascendc' could not load cann_ops_transformer. Install the standalone "
            "ops-transformer Python package with compatible PyTorch, torch-npu and CANN libraries. "
            f"See hyper_parallel/components/functional/README.md. Original error: {error}"
        ) from error
    block_attention_residuals_op = getattr(ops, "block_attention_residuals", None)
    block_attention_residuals_backward = getattr(ops, "block_attention_residuals_backward", None)
    if not callable(block_attention_residuals_op) or not callable(block_attention_residuals_backward):
        raise ImportError(
            "AttnRes backend='ascendc' requires cann_ops_transformer.block_attention_residuals and "
            "block_attention_residuals_backward. Install a package containing both training APIs "
            "and build block_attention_residuals plus block_attention_residuals_grad. "
            "See hyper_parallel/components/functional/README.md."

            "block_attention_residuals_backward. Install a package containing both training APIs "
            "and build block_attention_residuals plus block_attention_residuals_grad. "
            "See hyper_parallel/components/functional/README.md."
        )
    return block_attention_residuals_op


def _block_attention_residuals_fused(prefix: torch.Tensor, history: torch.Tensor,
                                    projection: torch.Tensor, norm_weight: torch.Tensor,
                                    norm_eps: float) -> torch.Tensor:
    """Prepare differentiable inputs and invoke the fused block attention residuals API."""
    block_attention_residuals_op = _load_block_attention_residuals()
    # Native forward accepts [H], but its training backward requires [1,H].
    if projection.ndim == 1:
        projection = projection.unsqueeze(0)
    inputs = (prefix, history, projection, norm_weight)
    common_dtype = prefix.dtype
    for tensor in inputs[1:]:
        common_dtype = torch.promote_types(common_dtype, tensor.dtype)
    # The native API requires one dtype; casts must preserve all four gradient edges.
    prepared = tuple(tensor.to(dtype=common_dtype).contiguous() for tensor in inputs)
    result = block_attention_residuals_op(*prepared, valid_block_num=-1, norm_eps=norm_eps)
    return result.to(torch.promote_types(prefix.dtype, history.dtype))


def block_attention_residuals(prefix: torch.Tensor, history: torch.Tensor,
                             projection: torch.Tensor, norm_weight: torch.Tensor, *, norm_eps: float,
                             backend: str = "eager") -> torch.Tensor:
    """Mix previous blocks and the current residual through depth attention.

        Models own block boundaries and learned parameters. The fused operator
        supports first-order gradients. Its dispatcher and native entry validate
        device and shape limits; kernel execution errors propagate unchanged.
    """
    if backend == "eager":
        return _block_attention_residuals_eager(prefix, history, projection, norm_weight, norm_eps)
    if backend != "ascendc":
        raise ValueError("AttnRes backend must be 'eager' or 'ascendc'")
    if prefix.shape[0] == 0 or history.shape[1] == 0:
        # Empty views keep zero gradients without inspecting unused values.
        zero = sum(tensor.reshape(-1)[:0].sum() for tensor in (history, projection, norm_weight))
        dtype = torch.promote_types(prefix.dtype, history.dtype)
        return prefix.to(dtype) + zero.to(dtype)
    return _block_attention_residuals_fused(prefix, history, projection, norm_weight, norm_eps)